diff --git a/.gitignore b/.gitignore index f5509855..b00a493e 100644 --- a/.gitignore +++ b/.gitignore @@ -20,7 +20,12 @@ logs *.log.* *.db *.lock +tmp +cache ### IDE .vscode -*.nogit.* \ No newline at end of file +*.nogit.* + +*.pdf +src/data \ No newline at end of file diff --git a/src/.env.template b/src/.env.template index c172cd7b..60e1b089 100644 --- a/src/.env.template +++ b/src/.env.template @@ -3,4 +3,6 @@ DEEPSEEKAPI=optional ZHIPUAPI=optional QIANFAN_ACCESS_KEY=optional QIANFAN_SECRET_KEY=optional -HF_ENDPOINT='https://hf-mirror.com' \ No newline at end of file +HF_ENDPOINT='https://hf-mirror.com' +CUDA_VISIBLE_DEVICES=0 +MODEL_ROOT_DIR=MODEL_ROOT_DIR_UNSET \ No newline at end of file diff --git a/src/cli.py b/src/cli.py index a342846f..1a0fd1bf 100644 --- a/src/cli.py +++ b/src/cli.py @@ -12,7 +12,6 @@ if __name__ == "__main__": config = Config("config/base.yaml") model = select_model(config) retriever = Retriever(config) - # pre_retrieval.add_file("/home/zwj/workspace/ProjectAthena/src/data/file/鉴定工作报告、技术报告-0708.pdf") print(f"[{config.model_provider}:{config.get('model_name', 'default')}] Type 'exit' to quit") diff --git a/src/config/__init__.py b/src/config/__init__.py index 94db249d..fe74ed63 100644 --- a/src/config/__init__.py +++ b/src/config/__init__.py @@ -45,6 +45,7 @@ class Config(SimpleConfig): ### handle local model model_root_dir = os.getenv("MODEL_ROOT_DIR", "pretrained_models") for model, model_rel_path in self.model_local_paths.items(): + # 如果 model_rel_path 不是绝对路径，那么拼接 model_root_dir if not model_rel_path.startswith("/"): self.model_local_paths[model] = os.path.join(model_root_dir, model_rel_path) diff --git a/src/config/base.yaml b/src/config/base.yaml index 233065be..866c6a54 100644 --- a/src/config/base.yaml +++ b/src/config/base.yaml @@ -3,6 +3,9 @@ name: base ## model ### model_provider, option in deepseek, zhipu model_provider: qianfan + +# 注意这里是模型名，而不是具体的模型路径，默认使用 HuggingFace 的路径 +# 如果需要自定义路径，则在下面配置 model_local_paths embed_model: bge-large-zh # option in ["bge-large-zh"] ## startup @@ -17,7 +20,8 @@ enable_search_engine: True ## vector_store milvus_local_path: data/vector_base/milvus.db # 初步 - -## model dir +## model dir 可以写相对路径和绝对路径 +### 相对路径是相对于环境变量中 MODEL_ROOT_DIR 的路径 model_local_paths: - bge-large-zh: bge-large-zh-v1.5 \ No newline at end of file + bge-large-zh: bge-large-zh-v1.5 + oneke: oneke \ No newline at end of file diff --git a/src/core/preretrieval.py b/src/core/preretrieval.py index b3251505..f10240c7 100644 --- a/src/core/preretrieval.py +++ b/src/core/preretrieval.py @@ -1,5 +1,6 @@ # Read Chunking Embedding and save it to Vector Database import os +import utils from pathlib import Path from llama_index.readers.file import PDFReader @@ -8,6 +9,7 @@ from models.embedding import EmbeddingModel from utils.logging_config import setup_logger from pymilvus import MilvusClient +from plugins import OneKE, pdf2txt logger = setup_logger("PreRetrieval") @@ -17,10 +19,12 @@ def pdfreader(file_path): assert os.path.exists(file_path), "File not found" assert file_path.endswith(".pdf"), "File format not supported" - doc = PDFReader().load_data(file=Path(file_path)) + if utils.is_text_pdf(file_path): + doc = PDFReader().load_data(file=Path(file_path)) + text = "\n\n".join([d.get_content() for d in doc]) + else: + text = pdf2txt(file_path) - # 简单的拼接起来之后返回纯文本 - text = "\n\n".join([d.get_content() for d in doc]) return text def plainreader(file_path): @@ -40,6 +44,7 @@ class PreRetrieval: self.embed_model = EmbeddingModel(config) self.client = MilvusClient(config.milvus_local_path) + self.oneke = OneKE(config) def _init_config(self, config): self.vector_dim = 1024 # 暂时不知道这个和 embedding model 的 embedding 大小有什么关系 @@ -50,6 +55,8 @@ class PreRetrieval: """添加文件到数据库""" collection_name = collection_name or self.default_collection_name text = self.read_text(file) + # convert text to graph + # self.oneke.processing_text_to_kg(text, none) chunks = self.chunking(text) self.add_documents(chunks, collection_name) diff --git a/src/models/embedding.py b/src/models/embedding.py index f4d7924a..a01c7c56 100644 --- a/src/models/embedding.py +++ b/src/models/embedding.py @@ -18,11 +18,7 @@ class EmbeddingModel(FlagModel): assert config.embed_model in SUPPORT_LIST.keys(), f"Unsupported embed model: {config.embed_model}, only support {SUPPORT_LIST.keys()}" - if config.embed_model in config.model_local_paths.keys(): - model_name_or_path = config.model_local_paths[config.embed_model] - else: - model_name_or_path = SUPPORT_LIST[config.embed_model] - + model_name_or_path = config.model_local_paths.get(config.embed_model, SUPPORT_LIST[config.embed_model]) logger.info(f"Loading embedding model {config.embed_model} from {model_name_or_path}") super().__init__(model_name_or_path, diff --git a/src/plugins/__init__.py b/src/plugins/__init__.py new file mode 100644 index 00000000..cae80c54 --- /dev/null +++ b/src/plugins/__init__.py @@ -0,0 +1,2 @@ +from plugins.oneke import * +from plugins.pdf2txt import * \ No newline at end of file diff --git a/src/plugins/oneke.py b/src/plugins/oneke.py index d6cb6ee0..0f1f3757 100644 --- a/src/plugins/oneke.py +++ b/src/plugins/oneke.py @@ -1,7 +1,8 @@ # readmore: https://github.com/zjunlp/DeepKE/blob/main/example/llm/OneKE.md - +import os import json import torch +import dotenv from transformers import ( AutoConfig, AutoTokenizer, @@ -10,11 +11,16 @@ from transformers import ( BitsAndBytesConfig ) -MODEL_NAME_OR_PATH = "zjunlp/OneKE" +from utils import setup_logger +logger = setup_logger("OneKE") + +dotenv.load_dotenv() + +MODEL_NAME_OR_PATH = os.path.join(os.getenv('MODEL_ROOT_DIR'), 'OneKE') instruction_mapper = { 'NERzh': "你是专门进行实体抽取的专家。请从input中抽取出符合schema定义的实体，不存在的实体类型返回空列表。请按照JSON字符串的格式回答。", - 'REzh': "你是专门进行关系抽取的专家。请从input中抽取出符合schema定义的关系三元组，不存在的关系返回空列表。请按照JSON字符串的格式回答。", + 'REzh': "你是专门进行关系抽取的专家。请从input中抽取出符合schema定义的关系三元组。请按照JSON字符串的格式回答。", 'EEzh': "你是专门进行事件提取的专家。请从input中抽取出符合schema定义的事件，不存在的事件返回空列表，不存在的论元返回NAN，如果论元存在多值请返回列表。请按照JSON字符串的格式回答。", 'EETzh': "你是专门进行事件提取的专家。请从input中抽取出符合schema定义的事件类型及事件触发词，不存在的事件返回空列表。请按照JSON字符串的格式回答。", 'EEAzh': "你是专门进行事件论元提取的专家。请从input中抽取出符合schema定义的事件论元及论元角色，不存在的论元返回NAN或空字典，如果论元存在多值请返回列表。请按照JSON字符串的格式回答。", @@ -33,11 +39,15 @@ split_num_mapper = { 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) + def __init__(self, config=None): + + self.config = config + model_name_or_path = config.model_local_paths.get('oneke', "zjunlp/OneKE") + logger.info(f"Loading KGC model OneKE from {model_name_or_path}") + + model_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) - # 4bit量化OneKE self.quantization_config=BitsAndBytesConfig( load_in_4bit=True, llm_int8_threshold=6.0, @@ -54,14 +64,13 @@ class OneKE: ) self.model = AutoModelForCausalLM.from_pretrained( - MODEL_NAME_OR_PATH, - config=self.config, + model_name_or_path, + config=model_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): @@ -89,10 +98,9 @@ class OneKE: 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") + input_ids = self.tokenizer.encode(sintruct, return_tensors="pt").to(self.model.device) input_length = input_ids.size(1) generation_output = self.model.generate( @@ -108,10 +116,137 @@ class OneKE: return outputs + def processing_text_to_kg(self, text_or_path, output_path): + for chunk in read_and_process_chars(text_or_path): + + text = chunk + schema = [ + { + "entity_type": "食品", + "attributes": { + "名称": "食品的名称，包括品牌名、通用名称或专业化学名", + "分类": "食品所属的类型，例如水果、蔬菜、肉类、谷物、调料、添加剂、益生菌等", + "成分": "食品的主要成分，详细列出包括天然成分、添加剂、保鲜剂、营养强化剂等", + "营养价值": "食品的营养成分，概括其提供的能量和主要营养素，如蛋白质、脂肪、碳水化合物、维生素和矿物质", + "加工方式": "食品的处理或制备方法，包括日常烹饪、加工处理及实验室制备方式等", + "作用或食用效果": "食品对健康或身体的影响，可能的功效或用途" + } + } + ] + task = "KG" + output = self.predict(text=text, schema=schema, task=task, language="zh") + formatted_output = parse_and_format_output(output=output, task_type=task) + + with open(output_path, 'a+', encoding='utf-8') as f: + for entry in formatted_output: + f.write(json.dumps(entry, ensure_ascii=False) + '\n') + + print(f"预测结果已添加到 {output_path} 文件中。") + return output_path + +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() - text = "《红楼梦》是一部中国古典长篇小说，是清代作家曹雪芹创作的作品。" - schema = ["人物", "地点", "时间", "事件"] - output = oneke.predict(text, schema, "NER") - print(output) \ No newline at end of file + oneke.processing_text_to_kg("asdasdadadsad", 'kg.jsonl') + + file_path = '' + output_path = '' + 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) + + with open(output_path, 'a+', encoding='utf-8') as f: + for entry in formatted_output: + f.write(json.dumps(entry, ensure_ascii=False) + '\n') + + print(f"预测结果已添加到 {output_path} 文件中。") \ No newline at end of file diff --git a/src/plugins/pdf2txt.py b/src/plugins/pdf2txt.py new file mode 100644 index 00000000..22c6f373 --- /dev/null +++ b/src/plugins/pdf2txt.py @@ -0,0 +1,70 @@ +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 + + +def pdf2txt(pdf_path): + output_dir = os.path.join('tmp', 'pdf2txt', os.path.basename(pdf_path).split('.')[0]) + os.makedirs(output_dir, exist_ok=True) + + table_engine = PPStructure(recovery=True, lang='ch') + + imgs = [] + img_dir = os.path.join(output_dir, 'imgs') + if not os.path.exists(img_dir): + os.makedirs(img_dir) + pdfDoc = fitz.open(pdf_path) + 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) + img_filename = os.path.join(img_dir, f'images_{pg+1}.png') + pix.save(img_filename) # os.sep + imgs.append(img_filename) + else: + img_names = sorted(os.listdir(img_dir)) + imgs = [os.path.join(img_dir, img_name) for img_name in img_names] + + respath = os.path.join(output_dir, 'res.txt') + + text = [] + for img_name in tqdm(imgs, desc='to txt', ncols=100): + img = cv2.imread(img_name) + result = table_engine(img) + + save_structure_res(result, output_dir, "structure_result") + + res = sorted_layout_boxes(result, img.shape[1]) + + try: + convert_info_docx(img, res, output_dir, "info") + 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') # okk + for pra in line['res']: + if isinstance(pra, dict) and 'text' in pra: + text.append(pra['text']) + else: + continue # 如果不是字典或者缺少 'text' 键，跳过当前循环 + text.append('\n') + + whole_text = '\n'.join(text) + with open(respath, 'w', encoding='utf-8') as f: + f.write(whole_text) + + return whole_text + +if __name__ == "__main__": + pdf_path = r'data/file/焙烤食品工艺学.pdf' + print(pdf2txt(pdf_path)) diff --git a/src/utils/__init__.py b/src/utils/__init__.py new file mode 100644 index 00000000..27224e83 --- /dev/null +++ b/src/utils/__init__.py @@ -0,0 +1,12 @@ +import fitz + +from utils.logging_config import setup_logger, logger + +def is_text_pdf(pdf_path): + doc = fitz.open(pdf_path) + for page_num in range(len(doc)): + page = doc.load_page(page_num) + text = page.get_text() + if text.strip(): # 检查是否有文本内容 + return True + return False \ No newline at end of file diff --git a/src/utils/logging_config.py b/src/utils/logging_config.py index 998a856b..0dd3d979 100644 --- a/src/utils/logging_config.py +++ b/src/utils/logging_config.py @@ -1,4 +1,5 @@ import logging +import os from datetime import datetime @@ -9,7 +10,7 @@ def setup_logger(name, log_file=None, level=logging.DEBUG, console=True): if log_file is None: log_file = f'output/log/project-{DATETIME}.log' - + os.makedirs("output/log", exist_ok=True) """Function to setup logger with the given name and log file.""" logger = logging.getLogger(name)