output/test_knn_2/one_random.py

# -*- coding: utf-8 -*-
import pandas as pd
import json

# 加载地质描述文本，提取prompt和label
with open('F:/GeoLLM/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']
# print(prompt,label)

print(len(text))
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 = 'gpt'
# model_name = 'deepseek-ai/DeepSeek-V3'
model_name='deepseek-ai/DeepSeek-R1'

prompt = '''
你是一名专业经验丰富的工程地质领域专家，你的任务是从给定的输入文本中提取"实体-关系-实体"三元组。关系类型包括24种："出露于"、"位于"、"整合接触"、"不整合接触"、"假整合接触"、"断层接触"、"分布形态"、"大地构造位置"、"地层区划"、"出露地层"、"岩性"、"厚度"、"面积"、"坐标"、"长度"、"含有"、"所属年代"、"行政区划"、"发育"、"古生物"、"海拔"、"属于"、"吞噬"、"侵入"。提取过程请按照以下规范：
1.	输出格式：
严格遵循JSON数组，无额外文本，每个元素包含：
[
    {
        "entity1": "实体1",
        "relation": "关系",
        "entity2": "实体2"
    }
]
2.	复杂关系处理：
- 若同一实体参与多个关系，需分别列出不同三元组
'''
j=0
q=0

# 统计目前处理到第几条
json_path = 'F:/GeoLLM/output/knn/two_shot/'+model_name+'.json'
j=len(json.load(open(json_path,'r',encoding='utf-8')))
q=len(json.load(open('F:/GeoLLM/output/knn/two_shot_raw/'+model_name+'.json','r',encoding='utf-8')))
print(j)
# # 当q=j时才继续处理
if q==j:        
    for i in range(j,500):
        # gemini模型是免费的，但是有请求限制，所以需要等待1秒，否则会被截断
        # 这似乎不起作用，因此采用和gpt相同的方式搭建

        # sleep(3)
        # 从text的500-1000数据中随机获取一个完整的text和triple_list作为提示
        # prompt_string = generate_prompt(text, label, 1)
        # 从text的500-1000数据中生成与text[i]最相似的3个样例的提示
        prompt_string_b = generate_prompt_b(text, label, text[i], 2)
        # print(prompt_string_b)
        # response = zero_shot(model_series, model_name, prompt+text[i])
        # print(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])
        print(i)
        # 解析响应
        formatted_triples = parse_llm_response(response)
        
        # 保存结果
        save_to_json(text[i], formatted_triples, model_series=model_name, output_dir='F:/GeoLLM/output/knn/two_shot/')
        # 保存原始响应为josn文件save_raw_response
        save_raw_response(response, text[i], model_series=model_name, output_dir='F:/GeoLLM/output/knn/two_shot_raw/')

else:
    print('q!=j')