GeoLLM / 0423_and_25_test /0423_gmini.py
Ciallo0d00's picture
Upload folder using huggingface_hub
badcf3c verified
# 测试llama和qen
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']
# print(prompt,label)
print(len(text))
from response_to_json import parse_llm_response, save_to_json, save_raw_response
from LLM_Special_for_0423 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 = 'gemini'
model_name = 'gemini-1.5-pro-002'
prompt = '''
你是一名专业经验丰富的工程地质领域专家,你的任务是从给定的输入文本中提取"实体-关系-实体"三元组。关系类型包括24种:"出露于"、"位于"、"整合接触"、"不整合接触"、"假整合接触"、"断层接触"、"分布形态"、"大地构造位置"、"地层区划"、"出露地层"、"岩性"、"厚度"、"面积"、"坐标"、"长度"、"含有"、"所属年代"、"行政区划"、"发育"、"古生物"、"海拔"、"属于"、"吞噬"、"侵入"。提取过程请按照以下规范:
1.  输出格式:
严格遵循JSON数组,无额外文本,每个元素包含:
[
    {
        "entity1": "实体1",
        "relation": "关系",
        "entity2": "实体2"
    }
]
2.  复杂关系处理:
若同一实体参与多个关系,需分别列出不同三元组
3.  关系的解释:
对部分较难提取的关系解释如下:
(1)假整合接触:特指产状一致的平行不整合类型,强调沉积序列的间断但无构造变形,专门用来描述地层之间的关系,请勿滥用。例如:“索拉克组与上部中奥陶统林组呈为平行不整合接触” 可提取出三元组(索拉克组,假整合接触,中奥陶统林组)。
(2)大地构造位置:定位地质单元在板块构造格架中的归属,关联造山带或构造单元划分。在三元组中主体和客体分别是"岩石、地层(X群或X系)"和"XX带"。例如:“石炭系的大地构造位置位于冈瓦纳构造带北部边缘。” 可提取出三元组(石炭系,大地构造位置,冈瓦纳构造带北缘)。
(3)出露地层:特指某区域实际暴露的地层实体,强调可观测的地表地质单元。在三元组中主体和客体分别是"地层(X组或X群)"和"地层(X组或X群)"。例如:“红柳沟金铜矿区出露的地层主要为南华 - 下奥陶统红柳沟群。” 可提取出三元组(红柳沟金铜矿区,出露地层,南华 - 下奥陶统红柳沟群)。
(4)坐标:特指记录地质特征点的地理空间定位数据。例如:“索拉克铜金矿中心的地理坐标为东经90°11′47″,北纬39°12′15″。” 可提取出三元组(索拉克铜金矿址,坐标,东经90°11′47″)。
(5)发育:描述地质构造或沉积特征的显现程度与形成状态强度。在三元组中主体和客体分别是"岩石"和"X层(片)理等、X构造"。例如:“兰花瓮组发育水平层理和条带。” 可提取出三元组(兰花瓮组,发育,水平层理)。
(6)古生物:记录地层中赋存的化石生物信息,需包含完整拉丁学名与分类特征。专指"X岩石、地区、地层(X组或X群"有"古生物"。例如:“该套地层的古生物为 Lumu et al。” 可提取出三元组(地层,古生物,Lumu et a)。
(7)属于:确立地质单元在分类体系中的类型归属,构建概念层级关系,不可滥用。例如:“矿区属于红柳沟 - 拉渠泉金、铜、铅、锌、铁多金属矿化亚区。” 可提取出三元组(矿区,属于,多金属矿化亚区)。
'''
j=0
q=0
json_path = './output/Knowledge-guided_rerun/nomal/'+model_name+'.json'
j=len(json.load(open(json_path,'r',encoding='utf-8')))
q=len(json.load(open('./output/Knowledge-guided_rerun/nomal_raw/'+model_name+'.json','r',encoding='utf-8')))
# # 当q=j时才继续处理
# # 当q=j时才继续处理
if q==j:
print(j)
for i in range(j,500):
response = zero_shot(model_series, model_name, prompt+'\n'+'以下是地质描述文本,请基于上述描述提取三元组:'+'\n'+text[i])
print(prompt+'\n'+'以下是地质描述文本,请基于上述描述提取三元组:'+'\n'+text[i])
# 解析响应
formatted_triples = parse_llm_response(response)
# 保存结果
save_to_json(text[i], formatted_triples, model_series=model_name, output_dir='./output/Knowledge-guided_rerun/nomal/')
# 保存原始响应为josn文件save_raw_response
save_raw_response(response, prompt+'\n'+'以下是地质描述文本,请基于上述描述提取三元组:'+'\n'+text[i], model_series=model_name, output_dir='./output/Knowledge-guided_rerun/nomal_raw/')
else:
print('q!=j')