from sentence_transformers import SentenceTransformer
import scipy.spatial
import pandas as pd
'''
短文本推荐 模型 sentence bert
1.文本向量 对 词向量 的好处
2，sentence bert 对 bert 的好处
余弦相似度常用于文本相似度计算
而欧几里得距离和曼哈顿距离常用于数值型数据的相似度计算
Jaccard相似度常用于集合的相似度计算
'''

embedder = SentenceTransformer('bert-base-chinese')

df = pd.read_excel('../datas/骨科.xlsx', sheet_name='Sheet1')
corpus = df['ask'].to_numpy()

corpus_embeddings = embedder.encode(corpus)

# 待查询的句子
queries = ['3月来，患者左髋部，左膝部及左下肢出现疼痛，症状持续，左膝关节：左膝关节内半月板后角、外侧半月板信号局部增高，考虑变性。关节腔少量积液；髌下脂肪垫内絮状高信号影；c上囊滑囊，腘肌滑囊信号增高。膝关节后方静脉曲张？左侧股骨头形态改变，双侧股骨头及左侧髋臼异常信号，考虑缺血坏死。双侧髋关节骨质增生。左侧髋关节周围软组织水肿；肌肉萎缩']
query_embeddings = embedder.encode(queries)
closest_n = 20
# 对于每个句子，使用余弦相似度查询最接近的个句子closest n = 2
for query, query_embedding in zip(queries, query_embeddings):
    distances = scipy.spatial.distance.cdist([query_embedding], corpus_embeddings, "cosine")[0]
    # 按照距离逆序
    results = zip(range(len(distances)), distances)
    results = sorted(results, key=lambda x: x[1])
    print("====================================")
    print("Query:", query)
    print("Result:Top 20 most similar sentences in corpus:")
    for idx, distance in results[0:closest_n]:
        print(corpus[idx].strip(), "(Score: %.4f)" % (1 - distance))