from transformers import AutoTokenizer, AutoModel
import torch


# Mean Pooling - Take attention mask into account for correct averaging
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0] # First element of model_output contains all token embeddings
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)


# Sentences we want sentence embeddings for
sentences = ["경제 전문가가 금리 인하에 대한 예측을 하고 있다.", "주식 시장에서 한 투자자가 주식을 매수한다."]

# Load model from HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained("upskyy/kf-deberta-multitask")
model = AutoModel.from_pretrained("upskyy/kf-deberta-multitask")

# Tokenize sentences
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')

# Compute token embeddings
with torch.no_grad():
    model_output = model(**encoded_input)

# Perform pooling. In this case, mean pooling.
sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])

print("Sentence embeddings:")
print(sentence_embeddings)

from sentence_transformers import SentenceTransformer, util
import numpy as np

# Sentence transformer model for financial domain
embedder = SentenceTransformer("upskyy/kf-deberta-multitask")

# Financial domain corpus
corpus = [
    "주식 시장에서 한 투자자가 주식을 매수한다.",
    "은행에서 예금을 만기로 인출하는 고객이 있다.",
    "금융 전문가가 새로운 투자 전략을 개발하고 있다.",
    "증권사에서 주식 포트폴리오를 관리하는 팀이 있다.",
    "금융 거래소에서 새로운 디지털 자산이 상장된다.",
    "투자 은행가가 고객에게 재무 계획을 제안하고 있다.",
    "금융 회사에서 신용평가 모델을 업데이트하고 있다.",
    "투자자들이 새로운 ICO에 참여하려고 하고 있다.",
    "경제 전문가가 금리 인상에 대한 예측을 내리고 있다.",
]

corpus_embeddings = embedder.encode(corpus, convert_to_tensor=True)

# Financial domain queries
queries = [
    "한 투자자가 비트코인을 매수한다.",
    "은행에서 대출을 상환하는 고객이 있다.",
    "금융 분야에서 새로운 기술 동향을 조사하고 있다."
]

# Find the closest 5 sentences of the corpus for each query sentence based on cosine similarity
top_k = 5
for query in queries:
    query_embedding = embedder.encode(query, convert_to_tensor=True)
    cos_scores = util.pytorch_cos_sim(query_embedding, corpus_embeddings)[0]
    cos_scores = cos_scores.cpu()

    # We use np.argpartition, to only partially sort the top_k results
    top_results = np.argpartition(-cos_scores, range(top_k))[0:top_k]

    print("\n\n======================\n\n")
    print("Query:", query)
    print("\nTop 5 most similar sentences in the financial corpus:")

    for idx in top_results[0:top_k]:
        print(corpus[idx].strip(), "(Score: %.4f)" % (cos_scores[idx]))