Spaces:
Runtime error
Runtime error
File size: 4,479 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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
import json
import os
def parse_llm_response(response):
"""
解析LLM响应并转换为标准格式
Args:
response: 可以是字符串或带有content属性的对象
"""
# 检查response类型并获取content
if isinstance(response, str):
content = response
elif isinstance(response, dict) and 'content' in response:
content = response['content']
else:
content = response.content
# 尝试查找JSON内容
try:
# 首先尝试查找markdown JSON块
if '```json' in content:
content = content.split('```json\n')[1].split('\n```')[0]
elif '```' in content:
content = content.split('```\n')[1].split('\n```')[0]
# 如果上述方法失败，尝试查找方括号包围的JSON数组
if content.find('[') != -1 and content.find(']') != -1:
start = content.find('[')
end = content.rfind(']') + 1
content = content[start:end]
# 清理可能的多余空白字符
content = content.strip()
# 解析JSON
triples = json.loads(content)
# # 在解析时,检查是否存在实体和关系,不存在报错，提示重新生成
# cleaned_triples = []
# for triple in triples:
# cleaned = {
# "entity1": triple["entity1"],
# "relation": triple["relation"],
# "entity2": triple["entity2"]
# }
# cleaned_triples.append(cleaned)
# 转换格式，检查是否存在实体和关系,不存在报错，提示重新生成
formatted_triples = {
"triple_list": [
[triple["entity1"], triple["relation"], triple["entity2"]]
for triple in triples
if all(key in triple for key in ["entity1", "relation", "entity2"])
]
}
# "triple_list": cleaned_triples
return formatted_triples
except json.JSONDecodeError as e:
print(f"JSON解析错误,原始内容: {content}")
print(f"错误详情: {str(e)}")
return None
def save_to_json(text, formatted_triples, model_series, output_dir='./output'):
"""
保存结果到JSON文件
"""
# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)
output_path = f'{output_dir}/{model_series}.json'
# 如果文件不存在则创建新文件,否则读取已有内容
if os.path.exists(output_path):
with open(output_path, 'r', encoding='utf-8') as f:
existing_data = json.load(f)
else:
existing_data = []
# 创建包含文本和三元组的新数据项
new_item = {
"text": text,
"triple_list": formatted_triples["triple_list"]
}
# 将新的数据项添加到已有数据中
if isinstance(existing_data, dict):
existing_data = [existing_data]
existing_data.append(new_item)
# 输出更新后的JSON格式
with open(output_path, 'w', encoding='utf-8') as f:
json.dump(existing_data, f, ensure_ascii=False, indent=4)
# 保存原始响应函数，仿照save_to_json
def save_raw_response(response, prompt, model_series, output_dir='./output/two_shot_raw'):
"""
保存原始响应到JSON文件
"""
# 确保输出目录存在
os.makedirs(output_dir, exist_ok=True)
output_path = f'{output_dir}/{model_series}.json'
# 如果文件不存在则创建新文件,否则读取已有内容
if os.path.exists(output_path):
with open(output_path, 'r', encoding='utf-8') as f:
existing_data = json.load(f)
else:
existing_data = []
# 创建包含文本和三元组的新数据项
new_item = {
"prompt": prompt,
# 解析response中的/n为回车
# AttributeError: 'dict' object has no attribute 'replace'
# 将response转换为字符串
"response": response.replace('\\n', '\n') if isinstance(response, str) else str(response)
}
# 将新的数据项添加到已有数据中
existing_data.append(new_item)
# 输出更新后的JSON格式
with open(output_path, 'w', encoding='utf-8') as f:
json.dump(existing_data, f, ensure_ascii=False, indent=4)
|