Merge pull request #2 from xerrors/batman

Batman

.gitignore

@@ -20,7 +20,12 @@ logs
 *.log.*
 *.db
 *.lock
+tmp
+cache
 
 ### IDE
 .vscode
 *.nogit.*
+
+*.pdf
+src/data

src/.env.template

@@ -4,3 +4,5 @@ ZHIPUAPI=optional
 QIANFAN_ACCESS_KEY=optional
 QIANFAN_SECRET_KEY=optional
 HF_ENDPOINT='https://hf-mirror.com'
+CUDA_VISIBLE_DEVICES=0
+MODEL_ROOT_DIR=MODEL_ROOT_DIR_UNSET

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")
 

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)
 

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
+  oneke: oneke

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"
 
+    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)
+
     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)

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,

src/plugins/__init__.py

@@ -0,0 +1,2 @@
+from plugins.oneke import *
+from plugins.pdf2txt import *

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)
+    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} 文件中。")

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))

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

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=False):
 
     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)