File size: 3,457 Bytes
badcf3c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
if __name__ == '__main__':
import pandas as pd
import json
# 加载地质描述文本，提取prompt和label
with open('./data/train_triples.json', 'r', encoding='utf-8') as f:
data = json.load(f)
# 将data转换为DataFrame
df = pd.DataFrame(data)
# 提取prompt和label
text = df['text']
label = df['triple_list']
from response_to_json import parse_llm_response, save_to_json, save_raw_response
from LLM import zero_shot
from prompt_generate import generate_prompt_with_examples as generate_prompt
from prompt_generate import generate_prompt_with_best_matches as generate_prompt_b
model_series = 'deepSeek'
# 提示 $0.5 / 1M tokens 补全 $1 / 1M tokens
# model_name = 'deepseek-ai/DeepSeek-V3'
# 提示 $1 / 1M tokens 补全 $4 / 1M tokens
# model_name='deepseek-ai/DeepSeek-R1'
# model_name = 'meta-llama/Meta-Llama-3.1-405B-Instruct'
model_name = 'Qwen/Qwen2.5-72B-Instruct'
prompt = '''
你是一名专业经验丰富的工程地质领域专家，你的任务是从给定的输入文本中提取"实体-关系-实体"三元组。关系类型包括24种："出露于"、"位于"、"整合接触"、"不整合接触"、"假整合接触"、"断层接触"、"分布形态"、"大地构造位置"、"地层区划"、"出露地层"、"岩性"、"厚度"、"面积"、"坐标"、"长度"、"含有"、"所属年代"、"行政区划"、"发育"、"古生物"、"海拔"、"属于"、"吞噬"、"侵入"。提取过程请按照以下规范：
1. 输出格式：
严格遵循JSON数组，无额外文本，每个元素包含：
[
{
"entity1": "实体1",
"relation": "关系",
"entity2": "实体2"
}
]
2. 复杂关系处理：
- 若同一实体参与多个关系，需分别列出不同三元组
'''
j=0
q=0
# json_path = './output/knn/three_shot/'+model_name+'.json'
# j=len(json.load(open(json_path,'r',encoding='utf-8')))
# q=len(json.load(open('./output/knn/three_shot_raw/'+model_name+'.json','r',encoding='utf-8')))
# # 当q=j时才继续处理
if q==j:
print(j)
for i in range(j,500):
# 从text的500-1000数据中随机获取一个完整的text和triple_list作为提示
# prompt_string = generate_prompt(text, label, 3)
# print(prompt_string)
prompt_string_b = generate_prompt_b(text, label, text[i], 3)
# response = zero_shot(model_series, model_name, prompt+text[i])
response = zero_shot(model_series, model_name, prompt+'\n'+'以下是地质描述文本和三元组提取样例'+'\n'+prompt_string_b+'\n'+'请根据样例提取三元组'+'\n'+text[i])
# print(prompt+'\n'+'以下是地质描述文本和三元组提取样例'+'\n'+prompt_string_b+'\n'+'请根据样例提取三元组'+'\n'+text[i])
# 解析响应
formatted_triples = parse_llm_response(response)
# 保存结果
save_to_json(text[i], formatted_triples, model_series=model_name, output_dir='./output/knn/three_shot/')
# 保存原始响应为josn文件save_raw_response
save_raw_response(response, text[i], model_series=model_name, output_dir='./output/knn/three_shot_raw/')
else:
print('q!=j') |