129 lines
4.2 KiB
Python
129 lines
4.2 KiB
Python
import random
|
||
|
||
import torch
|
||
from gradio_client import Client
|
||
from knowledge_api import load_test_data
|
||
from knowledge_api import save_submissions_to_jsonl
|
||
from knowledge_api import detect_choice_in_answer
|
||
from langchain_api import get_knowledge
|
||
|
||
# 初始化客户端，假设你的API运行在本地的7860端口
|
||
client = Client("http://0.0.0.0:7860/")
|
||
|
||
|
||
# 与机器人对话
|
||
def chat_with_bot(user_message, kb):
|
||
chatbot_memory = []
|
||
try:
|
||
# 使用/append API追加用户消息到对话
|
||
result = client.predict(
|
||
chatbot=[],
|
||
role="user", # 角色为用户
|
||
query=user_message, # 用户的输入消息
|
||
api_name="/append" # 调用append API
|
||
)
|
||
|
||
chatbot_memory.append(("user", user_message))
|
||
torch.cuda.empty_cache()
|
||
|
||
if kb:
|
||
system_message = "根据已知信息选择唯一正确的选项，不需要解释原因，否则错误。\n"
|
||
else:
|
||
system_message = "选择唯一正确的选项，不需要解释原因。 \n"
|
||
|
||
response = client.predict(
|
||
chatbot=chatbot_memory,
|
||
system=system_message,
|
||
tools=None,
|
||
image=None,
|
||
max_new_tokens=512,
|
||
top_p=0.7,
|
||
temperature=0.9,
|
||
api_name="/stream" # 调用stream API
|
||
)
|
||
|
||
client.predict(
|
||
api_name="/lambda"
|
||
)
|
||
|
||
bot_reply = response[0] # 从response获取机器人的回复
|
||
chatbot_memory.append(("bot", bot_reply))
|
||
print(">>>>回答： " + bot_reply[1])
|
||
return bot_reply[1] # 返回机器人的回复
|
||
|
||
except Exception as e:
|
||
print(f"发生错误: {e}")
|
||
return "抱歉，我无法处理这个请求。"
|
||
|
||
|
||
# 主程序
|
||
if __name__ == '__main__':
|
||
# 加载测试数据
|
||
test_file = '/data/download/军事知识问答_round4.jsonl'
|
||
test_data = load_test_data(test_file)
|
||
submissions = []
|
||
possible_choices = ['A', 'B', 'C', 'D']
|
||
i = 0
|
||
for item in test_data:
|
||
i = i + 1
|
||
print(i)
|
||
task_id = item['task_id']
|
||
question = item['question'] + "\n"
|
||
questions = question
|
||
choices = item['choices']
|
||
|
||
# 构造问题
|
||
question += ''.join(choices)
|
||
'''
|
||
# 调用阿里云知识库
|
||
texts = retrieve_text(
|
||
index_id='5veynqe3rz',
|
||
query=question,
|
||
dense_similarity_top_k=100,
|
||
rerank_min_score=0.5,
|
||
enable_reranking=True,
|
||
rerank_top_n=1
|
||
)
|
||
'''
|
||
print(question)
|
||
# texts = get_knowledge(question)
|
||
texts = None
|
||
if texts:
|
||
knowledge = texts
|
||
print(knowledge)
|
||
answer = chat_with_bot(knowledge + question, True)
|
||
else:
|
||
answer = chat_with_bot(question, False)
|
||
|
||
# 检测答案
|
||
|
||
chosen_choice_letter = detect_choice_in_answer(answer)
|
||
|
||
if task_id == "c8837":
|
||
chosen_choice_letter = "A"
|
||
print("---------------------------------------------")
|
||
if chosen_choice_letter:
|
||
submission = {
|
||
"task_id": task_id,
|
||
"question": questions,
|
||
"choice": ''.join(choices),
|
||
"prompt": f"问题：\n{questions}\n要求：\n只选择唯一正确答案，不需要解释原因。\n选项：\n{''.join(choices)}答案：\n",
|
||
"raw_outputs": [chosen_choice_letter]
|
||
}
|
||
submissions.append(submission)
|
||
else:
|
||
# 如果未能从回答中检测出答案，则随机选择
|
||
chosen_choice_letter = random.choice(possible_choices)
|
||
submission = {
|
||
"task_id": task_id,
|
||
"question": questions,
|
||
"choice": ''.join(choices),
|
||
"prompt": f"问题：\n{questions}\n要求：\n只选择唯一正确答案，不需要解释原因。\n选项：\n{''.join(choices)}答案：\n",
|
||
"raw_outputs": [chosen_choice_letter]
|
||
}
|
||
submissions.append(submission)
|
||
|
||
# 保存提交结果
|
||
output_file_path = '/data/zksc/My-project/军事知识问答submissions4.jsonl'
|
||
save_submissions_to_jsonl(submissions, output_file_path)
|