ocr and kg
This commit is contained in:
QiyiJiang
parent
401009bc9c
commit
f7132e321f
222
src/plugins/oneke copy.py
Normal file
222
src/plugins/oneke copy.py
Normal file
@ -0,0 +1,222 @@
|
||||
# readmore: https://github.com/zjunlp/DeepKE/blob/main/example/llm/OneKE.md
|
||||
import os
|
||||
import json
|
||||
import torch
|
||||
from transformers import (
|
||||
AutoConfig,
|
||||
AutoTokenizer,
|
||||
AutoModelForCausalLM,
|
||||
GenerationConfig,
|
||||
BitsAndBytesConfig
|
||||
)
|
||||
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
|
||||
MODEL_NAME_OR_PATH = "/data2024/yyyl/model/oneKE/"
|
||||
instruction_mapper = {
|
||||
'NERzh': "你是专门进行实体抽取的专家。请从input中抽取出符合schema定义的实体，不存在的实体类型返回空列表。请按照JSON字符串的格式回答。",
|
||||
'REzh': "你是专门进行关系抽取的专家。请从input中抽取出符合schema定义的关系三元组。请按照JSON字符串的格式回答。",
|
||||
'EEzh': "你是专门进行事件提取的专家。请从input中抽取出符合schema定义的事件，不存在的事件返回空列表，不存在的论元返回NAN，如果论元存在多值请返回列表。请按照JSON字符串的格式回答。",
|
||||
'EETzh': "你是专门进行事件提取的专家。请从input中抽取出符合schema定义的事件类型及事件触发词，不存在的事件返回空列表。请按照JSON字符串的格式回答。",
|
||||
'EEAzh': "你是专门进行事件论元提取的专家。请从input中抽取出符合schema定义的事件论元及论元角色，不存在的论元返回NAN或空字典，如果论元存在多值请返回列表。请按照JSON字符串的格式回答。",
|
||||
'KGzh': '你是一个图谱实体知识结构化专家。根据输入实体类型(entity type)的schema描述，从文本中抽取出相应的实体实例和其属性信息，不存在的属性不输出, 属性存在多值就返回列表，并输出为可解析的json格式。',
|
||||
'NERen': "You are an expert in named entity recognition. Please extract entities that match the schema definition from the input. Return an empty list if the entity type does not exist. Please respond in the format of a JSON string.",
|
||||
'REen': "You are an expert in relationship extraction. Please extract relationship triples that match the schema definition from the input. Return an empty list for relationships that do not exist. Please respond in the format of a JSON string.",
|
||||
'EEen': "You are an expert in event extraction. Please extract events from the input that conform to the schema definition. Return an empty list for events that do not exist, and return NAN for arguments that do not exist. If an argument has multiple values, please return a list. Respond in the format of a JSON string.",
|
||||
'EETen': "You are an expert in event extraction. Please extract event types and event trigger words from the input that conform to the schema definition. Return an empty list for non-existent events. Please respond in the format of a JSON string.",
|
||||
'EEAen': "You are an expert in event argument extraction. Please extract event arguments and their roles from the input that conform to the schema definition, which already includes event trigger words. If an argument does not exist, return NAN or an empty dictionary. Please respond in the format of a JSON string.",
|
||||
'KGen': 'You are an expert in structured knowledge systems for graph entities. Based on the schema description of the input entity type, you extract the corresponding entity instances and their attribute information from the text. Attributes that do not exist should not be output. If an attribute has multiple values, a list should be returned. The results should be output in a parsable JSON format.',
|
||||
}
|
||||
|
||||
split_num_mapper = {
|
||||
'NER':6, 'RE':4, 'EE':4, 'EET':4, 'EEA':4, 'KG':1
|
||||
}
|
||||
|
||||
class OneKE:
|
||||
|
||||
def __init__(self):
|
||||
self.config = AutoConfig.from_pretrained(MODEL_NAME_OR_PATH, trust_remote_code=True)
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_OR_PATH, trust_remote_code=True)
|
||||
|
||||
self.quantization_config=BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
llm_int8_threshold=6.0,
|
||||
llm_int8_has_fp16_weight=False,
|
||||
bnb_4bit_compute_dtype=torch.bfloat16,
|
||||
bnb_4bit_use_double_quant=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
)
|
||||
|
||||
self.generate_config = GenerationConfig(
|
||||
max_length=1024,
|
||||
max_new_tokens=512,
|
||||
return_dict_in_generate=True
|
||||
)
|
||||
|
||||
self.model = AutoModelForCausalLM.from_pretrained(
|
||||
MODEL_NAME_OR_PATH,
|
||||
config=self.config,
|
||||
device_map="auto",
|
||||
quantization_config=self.quantization_config,
|
||||
torch_dtype=torch.bfloat16,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
self.model.eval()
|
||||
|
||||
def construct_input(self, text, schema, task, language="zh", use_split=False):
|
||||
if use_split:
|
||||
split_num = split_num_mapper[task]
|
||||
if isinstance(schema, dict):
|
||||
schema_keys = list(schema.keys())
|
||||
schema_keys = [schema_keys[i:i+split_num] for i in range(0, len(schema_keys), split_num)]
|
||||
schema = [{k: schema[k] for k in keys} for keys in schema_keys]
|
||||
else:
|
||||
schema = [schema[i:i+split_num] for i in range(0, len(schema), split_num)]
|
||||
else:
|
||||
schema = [schema]
|
||||
|
||||
sintructs = []
|
||||
for s in schema:
|
||||
sintruct = json.dumps({
|
||||
"instruction": instruction_mapper[task+language],
|
||||
"schema": s,
|
||||
"input": text,
|
||||
}, ensure_ascii=False)
|
||||
sintructs.append(sintruct)
|
||||
|
||||
return sintructs
|
||||
|
||||
def predict(self, text, schema, task, language="zh", use_split=False):
|
||||
sintructs = self.construct_input(text, schema, task, language, use_split)
|
||||
outputs = []
|
||||
for sintruct in sintructs:
|
||||
input_ids = self.tokenizer.encode(sintruct, return_tensors="pt").to(self.model.device)
|
||||
input_length = input_ids.size(1)
|
||||
|
||||
generation_output = self.model.generate(
|
||||
input_ids=input_ids,
|
||||
generation_config=self.generate_config,
|
||||
pad_token_id=self.tokenizer.eos_token_id
|
||||
)
|
||||
generation_output = generation_output.sequences[0]
|
||||
generation_output = generation_output[input_length:]
|
||||
output = self.tokenizer.decode(generation_output, skip_special_tokens=True)
|
||||
|
||||
outputs.append(output)
|
||||
|
||||
return outputs
|
||||
|
||||
def read_and_process_chars(file_path, char_size=512, overlap_size=100):
|
||||
buffer = ""
|
||||
with open(file_path, 'r', encoding='utf-8') as file:
|
||||
while True:
|
||||
chunk = file.read(char_size)
|
||||
if not chunk: # 文件读取完毕
|
||||
if buffer:
|
||||
yield buffer
|
||||
break
|
||||
chunk = chunk.replace('\n', '').replace('\r', '') # 去除换行符
|
||||
buffer += chunk
|
||||
while len(buffer) >= char_size:
|
||||
yield buffer[:char_size]
|
||||
buffer = buffer[char_size - overlap_size:]
|
||||
|
||||
def parse_and_format_output(output, task_type):
|
||||
formatted_output = []
|
||||
|
||||
for entry in output:
|
||||
try:
|
||||
# Check if the entry is a valid JSON string
|
||||
if isinstance(entry, str) and entry.strip().startswith('{') and entry.strip().endswith('}'):
|
||||
parsed_entry = json.loads(entry)
|
||||
|
||||
if task_type == "KG":
|
||||
for entity_type, entities in parsed_entry.items():
|
||||
for entity_name, attributes in entities.items():
|
||||
for attribute_name, attribute_values in attributes.items():
|
||||
if isinstance(attribute_values, list):
|
||||
for value in attribute_values:
|
||||
formatted_output.append({
|
||||
"h": entity_name,
|
||||
"t": value,
|
||||
"r": attribute_name
|
||||
})
|
||||
else:
|
||||
formatted_output.append({
|
||||
"h": entity_name,
|
||||
"t": attribute_values,
|
||||
"r": attribute_name
|
||||
})
|
||||
elif task_type == "RE":
|
||||
for relation_type, pairs in parsed_entry.items():
|
||||
for pair in pairs:
|
||||
formatted_output.append({
|
||||
"h": pair["subject"],
|
||||
"t": pair["object"],
|
||||
"r": relation_type
|
||||
})
|
||||
else:
|
||||
raise json.JSONDecodeError("Invalid JSON format", entry, 0)
|
||||
except json.JSONDecodeError as e:
|
||||
print(f"JSONDecodeError: {e} - Skipping entry")
|
||||
continue
|
||||
except TypeError as e:
|
||||
print(f"TypeError: {e} - Skipping entry")
|
||||
continue
|
||||
except AttributeError as e:
|
||||
print(f"AttributeError: {e} - Skipping entry")
|
||||
continue
|
||||
|
||||
return formatted_output
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
oneke = OneKE()
|
||||
|
||||
file_path = 'food.txt'
|
||||
text = ""
|
||||
task = "KG"
|
||||
for chunk in read_and_process_chars(file_path):
|
||||
|
||||
text = chunk
|
||||
|
||||
# schema = {
|
||||
# "定义": "描述食品的起源、传统制作方法、文化象征意义或者描述食品或相关事物的定义或含义。包括食品的来源、特点及其在特定文化或背景下的意义。",
|
||||
# "组成成分": "描述食品的组成部分或成分，包括主要成分、微量成分、添加剂等，揭示食品的化学或物理组成。",
|
||||
# "功能": "描述食品或其成分的功能或作用，包括其对人体健康的影响、在烹饪中的用途、药用价值等。",
|
||||
# "属性": "描述食品或其成分的特性或属性，揭示其独特的营养价值、口感特点、保存方式等，包括食品的营养价值、口感特点（如口感丰富、清淡等）、保存方式（如冷藏、冷冻、干燥等）",
|
||||
# "种类": "描述食品或相关事物的种类或类别，揭示其分类体系、不同类型的特点及其在具体应用中的区别。"
|
||||
# }
|
||||
|
||||
schema = [
|
||||
{
|
||||
"entity_type": "食品",
|
||||
"attributes": {
|
||||
"名称": "食品的名称，包括品牌名、通用名称或专业化学名",
|
||||
"分类": "食品所属的类型，例如水果、蔬菜、肉类、谷物、调料、添加剂、益生菌等",
|
||||
"成分": "食品的主要成分，详细列出包括天然成分、添加剂、保鲜剂、营养强化剂等",
|
||||
"营养价值": "食品的营养成分，概括其提供的能量和主要营养素，如蛋白质、脂肪、碳水化合物、维生素和矿物质",
|
||||
"加工方式": "食品的处理或制备方法，包括日常烹饪、加工处理及实验室制备方式等",
|
||||
"作用或食用效果": "食品对健康或身体的影响，可能的功效或用途"
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
output = oneke.predict(text=text, schema=schema, task=task, language="zh")
|
||||
|
||||
formatted_output = parse_and_format_output(output=output, task_type=task)
|
||||
|
||||
if os.path.exists('output.json'):
|
||||
with open('output.json', 'r', encoding='utf-8') as f:
|
||||
existing_data = json.load(f)
|
||||
else:
|
||||
existing_data = []
|
||||
|
||||
# 将新的数据添加到现有数据之后
|
||||
existing_data.extend(formatted_output)
|
||||
|
||||
# 将合并后的数据写回 JSON 文件
|
||||
with open('output.json', 'w', encoding='utf-8') as f:
|
||||
json.dump(existing_data, f, ensure_ascii=False, indent=4)
|
||||
|
||||
text = ""
|
||||
|
||||
print("预测结果已添加到 output.json 文件中。")
|
||||
62
src/plugins/pdf2txt.py
Normal file
62
src/plugins/pdf2txt.py
Normal file
@ -0,0 +1,62 @@
|
||||
import os
|
||||
import fitz # fitz就是pip install PyMuPDF
|
||||
import cv2
|
||||
from paddleocr import PPStructure, save_structure_res
|
||||
from paddleocr.ppstructure.recovery.recovery_to_doc import sorted_layout_boxes, convert_info_docx
|
||||
from copy import deepcopy
|
||||
from tqdm import tqdm
|
||||
|
||||
table_engine = PPStructure(recovery=True, lang='ch')
|
||||
|
||||
def pdf2txt(pdfPath):
|
||||
imagePath = os.path.join('./imgs', os.path.basename(pdfPath).split('.')[0])
|
||||
if not os.path.exists(imagePath):
|
||||
os.makedirs(imagePath)
|
||||
pdfDoc = fitz.open(pdfPath)
|
||||
totalPage = pdfDoc.page_count
|
||||
for pg in tqdm(range(totalPage), desc='to imgs', ncols=100):
|
||||
page = pdfDoc[pg]
|
||||
rotate = int(0)
|
||||
zoom_x = 2
|
||||
zoom_y = 2
|
||||
mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
|
||||
pix = page.get_pixmap(matrix=mat, alpha=False)
|
||||
pix.save(imagePath + '/' + f'images_{pg+1}.png')
|
||||
|
||||
img_path = imagePath
|
||||
textPath = os.path.join('./txt', os.path.basename(pdfPath).split('.')[0])
|
||||
if not os.path.exists(textPath):
|
||||
os.makedirs(textPath)
|
||||
respath = os.path.join(textPath, 'res.txt')
|
||||
text = []
|
||||
imgs = sorted(os.listdir(img_path))
|
||||
for img_name in tqdm(imgs, desc='to txt', ncols=100):
|
||||
img = cv2.imread(os.path.join(img_path, img_name))
|
||||
result = table_engine(img)
|
||||
|
||||
save_structure_res(result, textPath, os.path.basename(img_path).split('.')[0])
|
||||
|
||||
h, w, _ = img.shape
|
||||
res = sorted_layout_boxes(result, w)
|
||||
|
||||
try:
|
||||
convert_info_docx(img, res, textPath, os.path.basename(img_path).split('.')[0])
|
||||
except Exception as e:
|
||||
print(f"Exception occurred while converting to DOCX: {e}")
|
||||
continue # Skip this image if conversion fails
|
||||
|
||||
for line in res:
|
||||
line.pop('img')
|
||||
for pra in line['res']:
|
||||
if isinstance(pra, dict) and 'text' in pra:
|
||||
text.append(pra['text'])
|
||||
else:
|
||||
continue # 如果不是字典或者缺少 'text' 键，跳过当前循环
|
||||
text.append('\n')
|
||||
|
||||
with open(respath, 'w', encoding='utf-8') as f:
|
||||
f.write('\n'.join(text))
|
||||
|
||||
if __name__ == "__main__":
|
||||
pdfPath = r'ppp/焙烤食品工艺学.pdf'
|
||||
pdf2txt(pdfPath)
|
||||
Loading…
Reference in New Issue
Block a user