# 测试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 = 'llama'
    model_name = 'meta-llama/Meta-Llama-3.1-405B-Instruct'

    prompt = '''
    你是一名专业经验丰富的工程地质领域专家，你的任务是从给定的输入文本中提取"实体-关系-实体"三元组。关系类型包括24种："出露于"、"位于"、"整合接触"、"不整合接触"、"假整合接触"、"断层接触"、"分布形态"、"大地构造位置"、"地层区划"、"出露地层"、"岩性"、"厚度"、"面积"、"坐标"、"长度"、"含有"、"所属年代"、"行政区划"、"发育"、"古生物"、"海拔"、"属于"、"吞噬"、"侵入"。提取过程请按照以下规范：
    1.  输出格式：
        严格遵循JSON数组，无额外文本，每个元素包含：
        [
            {
                "entity1": "实体1",
                "relation": "关系",
                "entity2": "实体2"
            }
        ]
    2.  复杂关系处理：
        若同一实体参与多个关系，需分别列出不同三元组
    3.  关系的解释：
        （1）出露于：指岩石或地层暴露在地表或近地表，未被覆盖或埋藏，在三元组中常显示为"出露"、"分布"，主体和客体分别是"岩石、地层（X组或X群）"和"地区"。例如：“晚奥陶世 - 志留世侵入岩主要出露于调查区南部。” 可以提取出三元组（晚奥陶世 - 志留世侵入岩，出露于，调查区南部）。
        （2）位于：确立地质单元在更大空间框架（行政区域/构造单元）中的从属关系，表达宏观空间归属，不能用来描述"X岩"在"X地方"，使用"出露于"关系替代。例如：“库穆奇志留世基性岩石（Sνβμ）主要分布在红柳沟 - 拉渠泉南华 - 早古生代结合区，该区位于调查区中西部。” 可提取出三元组（库穆奇志留世玄武岩基性岩石，位于，调查区中西部）
        （3）整合接触：指示上下地层连续沉积形成的接触关系，反映无显著沉积间断的岩性渐变特征，专门用来描述地层之间的关系，请勿滥用。例如：“索拉克组与上部中奥陶统林组呈整合接触。” 可提取出三元组（索拉克组，整合接触，中奥陶统林组）。
        （4）不整合接触：描述存在沉积缺失的地层接触界面，包含角度差异或岩性突变的接触特征，专门用来描述地层之间的关系，请勿滥用。例如：“通子岩组与毛口组呈不整合接触。” 可提取出三元组（通子岩组，不整合接触，毛口组）。
        （5）假整合接触：特指产状一致的平行不整合类型，强调沉积序列的间断但无构造变形，专门用来描述地层之间的关系，请勿滥用。例如：“索拉克组与上部中奥陶统林组呈为平行不整合接触” 可提取出三元组（索拉克组，假整合接触，中奥陶统林组）。
        （6）断层接触：两个地层之间为断层带或断层面，往往伴随动力破碎等构造现象，专门用来描述地层之间的关系，请勿滥用。例如：“索拉克组与中 - 上奥陶统拉排泉组呈断层接触。” 可提取出三元组（索拉克组，断层接触，上奥陶统拉排泉组）。
        （7）分布形态：刻画地质单元的空间展布特征，包括几何形态与延伸方向的组合描述。在三元组中主体和客体分别是"岩石、地层（X组或X系）"和"X块、状、体、结构、展布"等。例如：“石炭系呈东西向带状分布。” 可提取出三元组（石炭系，分布形态，带状）。
        （8）大地构造位置：定位地质单元在板块构造格架中的归属，关联造山带或构造单元划分。在三元组中主体和客体分别是"岩石、地层（X群或X系）"和"XX带"。例如：“石炭系的大地构造位置位于冈瓦纳构造带北部边缘。” 可提取出三元组（石炭系，大地构造位置，冈瓦纳构造带北缘）。
        （9）地层区划：表征地层单元在区域地层划分体系中的层级归属与分区属性。在三元组中主体和客体分别是"地层、测区"和"X分区"。例如：“石炭系地层区划属于冈瓦纳地层分区。” 可提取出三元组（石炭系，地层区划，冈瓦纳）。
        （10）出露地层：特指某区域实际暴露的地层实体，强调可观测的地表地质单元。在三元组中主体和客体分别是"地层（X组或X群）"和"地层（X组或X群）"。例如：“红柳沟金铜矿区出露的地层主要为南华 - 下奥陶统红柳沟群。” 可提取出三元组（红柳沟金铜矿区，出露地层，南华 - 下奥陶统红柳沟群）。
        （11）岩性：定义岩石的物质组成与结构特征，包含复合岩性的层级描述要素。例如：“晚奥陶世 - 志留世正长岩岩体的主要岩性为蚀变正长岩。” 可提取出三元组（晚奥陶世 - 志留世正长岩，岩性，蚀变正长岩）。
        （12）厚度：量化地层/岩体的垂向尺度，包含绝对数值与相对描述的量纲表达。例如：“正长岩厚度大于35.60 m。” 可提取出三元组（正长岩，厚度，35.60 m）。
        （13）面积：表征地质单元的水平展布范围，以数值与单位组合的标准化形式呈现。例如：“侵入岩体在矿区的出露面积约为 20%，约为 54 m2。” 可提取出三元组（侵入岩，出露面积，54 m2）
        （14）坐标：特指记录地质特征点的地理空间定位数据。例如：“索拉克铜金矿中心的地理坐标为东经90°11′47″，北纬39°12′15″。” 可提取出三元组（索拉克铜金矿址，坐标，东经90°11′47″）。
        （15）长度：描述线性地质体的空间延伸尺度。例如：“什比恩断裂带（F2）位于调查区东南角，位于阿拉坦达巴 - 达坂境内，总长度约为 20m。” 可提取出三元组（什比恩断裂带，长度，20m）。
        （16）含有：指示主体物质的成分包含关系，专指矿物组成或化石赋存状态，不同于日常的意思，不能用来描述“X岩"含有"古生物"，要使用"古生物"关系代替。例如：“中灰黑色块状燧石含有燧石条带。” 可提取出三元组（中灰黑色块状燧石，含有，燧石条带）。
        （17）所属年代：建立地质单元与标准地质年代体系的对应关系。例如：“红柳沟金铜矿区所属年代为早 - 中二叠世。” 可提取出三元组（红柳沟金铜矿区，所属年代，早 - 中二叠世）。
        （18）行政区划：界定地质实体在行政管理体系中的隶属层级与地域归属。例如：“调查区行政区划属于西藏自治区丁青县、乌齐县、洛隆县、昌都县、江达县、贡觉县、茶阳县。” 可提取出三元组（调查区，行政区划，茶阳县）。
        （19）发育：描述地质构造或沉积特征的显现程度与形成状态强度。在三元组中主体和客体分别是"岩石"和"X层（片）理等、X构造"。例如：“兰花瓮组发育水平层理和条带。” 可提取出三元组（兰花瓮组，发育，水平层理）。
        （20）古生物：记录地层中赋存的化石生物信息，需包含完整拉丁学名与分类特征。专指"X岩石、地区、地层（X组或X群"有"古生物"。例如：“该套地层的古生物为 Lumu et al。” 可提取出三元组（地层，古生物，Lumu et a）。
        （21）海拔：量化地质特征点相对于海平面的高程数据。例如：“索拉克铜金矿平均海拔为海平面以上 2800m。” 可提取出三元组（索拉克铜金矿址，海拔，2800m）。
        （22）属于：确立地质单元在分类体系中的类型归属，构建概念层级关系，不可滥用。例如：“矿区属于红柳沟 - 拉渠泉金、铜、铅、锌、铁多金属矿化亚区。” 可提取出三元组（矿区，属于，多金属矿化亚区）。
        （23）吞噬：表征侵入体对围岩的空间取代过程，反映岩浆活动的改造作用。例如：“任天堂岩组的南侧被侏罗纪花岗岩所吞噬。” 可提取出三元组（任天堂岩组，吞噬，侏罗纪花岗岩）。
        （24）侵入：描述岩浆岩体贯入围岩的地质作用过程，包含接触变质等伴生现象。例如：“高州壳石组合的岩石被片麻状花岗岩侵入。” 可提取出三元组（高州壳石组，侵入，片麻状花岗岩）。
    '''
    j=0
    q=0

    json_path = './output/Knowledge-guided_rerun/nomal_all/'+model_name+'.json'
    j=len(json.load(open(json_path,'r',encoding='utf-8')))
    q=len(json.load(open('./output/Knowledge-guided_rerun/nomal_all_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_all/')
            # 保存原始响应为josn文件save_raw_response
            save_raw_response(response, prompt+'\n'+'以下是地质描述文本，请基于上述描述提取三元组：'+'\n'+text[i], model_series=model_name, output_dir='./output/Knowledge-guided_rerun/nomal_all_raw/')

    else:
        print('q!=j')


# 上一轮提示 0421
    # 你是一名专业经验丰富的工程地质领域专家，你的任务是从给定的输入文本中提取"实体-关系-实体"三元组。关系类型包括24种："出露于"、"位于"、"整合接触"、"不整合接触"、"假整合接触"、"断层接触"、"分布形态"、"大地构造位置"、"地层区划"、"出露地层"、"岩性"、"厚度"、"面积"、"坐标"、"长度"、"含有"、"所属年代"、"行政区划"、"发育"、"古生物"、"海拔"、"属于"、"吞噬"、"侵入"。提取过程请按照以下规范：
    # 1.	输出格式：
    # 严格遵循JSON数组，无额外文本，每个元素包含：
    # [
    #     {
    #         "entity1": "实体1",
    #         "relation": "关系",
    #         "entity2": "实体2"
    #     }
    # ]
    # 2.	复杂关系处理：
    # - 若同一实体参与多个关系，需分别列出不同三元组
    # 3.	关系的解释：
    # 出露于：指岩石或地层暴露在地表或近地表，未被覆盖或埋藏，在三元组中常显示为"出露"、"分布"，主体和客体分别是"岩石、地层（X组或X群）"和"地区"。

    # 位于：确立地质单元在更大空间框架（行政区域/构造单元）中的从属关系，表达宏观空间归属，不能用来描述"X岩"在"X地方"，使用"出露于"关系替代。

    # 整合接触：指示上下地层连续沉积形成的接触关系，反映无显著沉积间断的岩性渐变特征，专门用来描述地层之间的关系，请勿滥用。

    # 不整合接触：描述存在沉积缺失的地层接触界面，包含角度差异或岩性突变的接触特征，专门用来描述地层之间的关系，请勿滥用。

    # 假整合接触：特指产状一致的平行不整合类型，强调沉积序列的间断但无构造变形，专门用来描述地层之间的关系，请勿滥用。

    # 断层接触：界定以断裂构造为边界的地质体空间关系，伴随破碎带等构造现象，专门用来描述地层之间的关系，请勿滥用。

    # 分布形态：刻画地质单元的空间展布特征，包括几何形态与延伸方向的组合描述。在三元组中主体和客体分别是"岩石、地层（X组或X群）"和"块、状、体、结构、展布"等。

    # 大地构造位置：定位地质单元在板块构造格架中的归属，关联造山带或构造单元划分。在三元组中主体和客体分别是"岩石、地层（X组或X群）"和"XX带"。

    # 地层区划：表征地层单元在区域地层划分体系中的层级归属与分区属性。在三元组中主体和客体分别是"地层、测区"和"X分区"。

    # 出露地层：特指某区域实际暴露的地层实体，强调可观测的地表地质单元。在三元组中主体和客体分别是"地层（X组或X群）"和"地层（X组或X群）"。

    # 岩性：定义岩石的物质组成与结构特征，包含复合岩性的层级描述要素。

    # 厚度：量化地层/岩体的垂向尺度，包含绝对数值与相对描述的量纲表达。

    # 面积：表征地质单元的水平展布范围，以数值与单位组合的标准化形式呈现。

    # 坐标：特指记录地质特征点的地理空间定位数据，保持原始经纬度格式特征。

    # 长度：描述线性地质体的空间延伸尺度。

    # 含有：指示主体物质的成分包含关系，专指矿物组成或化石赋存状态，不同于日常的意思，不能用来描述“X岩"含有"古生物"，要使用"古生物"关系代替。

    # 所属年代：建立地质单元与标准地质年代体系的对应关系，采用国际年代术语，不可滥用。

    # 行政区划：界定地质实体在行政管理体系中的隶属层级与地域归属。

    # 发育：描述地质构造或沉积特征的显现程度与形成状态强度。在三元组中主体和客体分别是"岩石"和"X层（片）理等、X构造"。

    # 古生物：记录地层中赋存的化石生物信息，需包含完整拉丁学名与分类特征。专指"X岩石、地区、地层（X组或X群"有"古生物"。

    # 海拔：量化地质特征点相对于海平面的高程数据，保留测量基准标识。

    # 属于：确立地质单元在分类体系中的类型归属，构建概念层级关系，不可滥用。

    # 吞噬：表征侵入体对围岩的空间取代过程，反映岩浆活动的改造作用。

    # 侵入：描述岩浆岩体贯入围岩的地质作用过程，包含接触变质等伴生现象

    # 4.	其他要点：

    # 所有三元组关系必须是以上24种之一
    # 关系的实体不能是动词或者介词等无意义的词。且描述岩石和地层等实体等要根据原文尽可能完整，如"X带X岩"或"X世X岩","X层"或"X组"
    # 描述数字时，不必带上"大于"等模糊词