# 型预测
from sentence_transformers import SentenceTransformer
import scipy.spatial
import pandas as pd

embedder = SentenceTransformer('bert-base-chinese')
# corpus = ['女宝宝，刚7岁，这一年，察觉到，我家孩子身上肉很多，而且，食量非常的大，平时都不喜欢吃去玩，请问：小儿肥胖超重该如何治疗。',
#           '男孩子，刚4岁，最近，发现，我家孩子体重要比别的孩子重很多，而且，最近越来越能吃了，还特别的懒，请问：小儿肥胖超重该怎样医治。',
#           '男宝，已经5岁，今年，察觉到，孩子身上越来越肉乎了，同时，吃的饭也比一般孩子多，平时都不喜欢吃去玩，请问：小儿肥胖能吃该如何治疗。']

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

corpus_embeddings = embedder.encode(corpus)

# 待查询的句子
queries = ['从大约2岁起，我发现孩子的左脚大拇指向食指方面倾斜，平时走路作习都正常，没有不舒服的感觉，但是拇指后的关节非常大，相对右侧比较突出。']
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 5 most similar sentences in corpus:")
    for idx, distance in results[0:closest_n]:
        print(corpus[idx].strip(), "(Score: %.4f)" % (1 - distance))