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 bd06628c..4add6e9f 100644 --- a/src/config/__init__.py +++ b/src/config/__init__.py @@ -44,6 +44,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..87728b40 --- /dev/null +++ b/src/plugins/__init__.py @@ -0,0 +1,2 @@ +from oneke import * +from pdf2txt import * \ No newline at end of file diff --git a/src/plugins/oneke copy.py b/src/plugins/oneke copy.py deleted file mode 100644 index 4173cc8b..00000000 --- a/src/plugins/oneke copy.py +++ /dev/null @@ -1,222 +0,0 @@ -# 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 文件中。") \ No newline at end of file diff --git a/src/plugins/oneke.py b/src/plugins/oneke.py index d6cb6ee0..6d2bab75 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,141 @@ class OneKE: return outputs + def processing_text_to_kg(self, text_or_path, output_path): + if os.path.isfile(text_or_path): + with open(text_or_path, 'r', encoding='utf-8') as f: + text = f.read() + else: + text = text_or_path + + 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 index 5edfb99f..22c6f373 100644 --- a/src/plugins/pdf2txt.py +++ b/src/plugins/pdf2txt.py @@ -6,13 +6,18 @@ from paddleocr.ppstructure.recovery.recovery_to_doc import sorted_layout_boxes, 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) +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] @@ -21,32 +26,32 @@ def pdf2txt(pdfPath): 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_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') - 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)) + img = cv2.imread(img_name) result = table_engine(img) - save_structure_res(result, textPath, os.path.basename(img_path).split('.')[0]) + save_structure_res(result, output_dir, "structure_result") - h, w, _ = img.shape - res = sorted_layout_boxes(result, w) + res = sorted_layout_boxes(result, img.shape[1]) try: - convert_info_docx(img, res, textPath, os.path.basename(img_path).split('.')[0]) + 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') + line.pop('img') # okk for pra in line['res']: if isinstance(pra, dict) and 'text' in pra: text.append(pra['text']) @@ -54,9 +59,12 @@ def pdf2txt(pdfPath): continue # 如果不是字典或者缺少 'text' 键，跳过当前循环 text.append('\n') + whole_text = '\n'.join(text) with open(respath, 'w', encoding='utf-8') as f: - f.write('\n'.join(text)) + f.write(whole_text) + + return whole_text if __name__ == "__main__": - pdfPath = r'ppp/焙烤食品工艺学.pdf' - pdf2txt(pdfPath) + 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..25282eb9 --- /dev/null +++ b/src/utils/__init__.py @@ -0,0 +1,12 @@ +import fitz + +from 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