from sentence_transformers import SentenceTransformer as ST
import json
from query import input_query
from sBERT import multi_ans
from ErnieBot_turbo import ErnieBot
from GPT_turbo import GPT

def llm_model(name):
    if name=='ErnieBot':
        return ErnieBot
    elif name=='gpt':
        return GPT

with open(".\config.json", "r", encoding="utf-8") as f:
    config = json.load(f)
with open(".\dataset.json", "r", encoding="utf-8") as f:
    dataset = json.load(f)

model_embed = ST(config['model_embed'])
model_llm = llm_model(config['model_llm'])
embed_path, data_path = config['embedding_path'], config['data_path']
language = config['language']
while(True):
    json_path, query = input_query(embed_path, language, model_llm)
    #如果要比较和数据集答案的相似度，使用以下代码；否则将这部分注释
    #query=dataset['Equestion30']
    #json_path=".\\embedding\\E.json"
    #answer=dataset['answer30']
    #answer_embeddings=model_embed.encode([answer])
    #print("问题：",query)
    #如果要比较和数据集答案的相似度，使用以上代码；否则将这部分注释
    ##############################################
    with open(json_path,'r',encoding="utf-8") as f:
        dict = json.load(f)
    file_name, sentence_embedding, url=  [], [], []
    for i in dict:
        url.append(dict[i][0])
        file_name.append(i)
        sentence_embedding.append(dict[i][1])
    multi_ans(query, url, file_name, sentence_embedding, model_embed, model_llm, data_path, language)
    break