Update main.py

main.py

@@ -6,7 +6,10 @@ from pathlib import Path
 from typing import List, Dict, Optional
 from dataclasses import dataclass, asdict
 
-from mineru import Mineru, Layout, Table
+from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
+from magic_pdf.data.dataset import PymuDocDataset
+from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
+from magic_pdf.config.enums import SupportedPdfParseMethod
 from sentence_transformers import SentenceTransformer
 from llama_cpp import Llama
 from fastapi.encoders import jsonable_encoder
@@ -27,11 +30,10 @@ class ProductSpec:
 
 class PDFProcessor:
     def __init__(self):
-        self.mineru = Mineru()
         self.emb_model = SentenceTransformer('all-MiniLM-L6-v2')
-        
-        # Initialize LLM with automatic download
         self.llm = self._initialize_llm()
+        self.output_dir = Path("./output")
+        self.output_dir.mkdir(exist_ok=True)
         
     def _initialize_llm(self):
         """Initialize LLM with automatic download if needed"""
@@ -44,21 +46,89 @@ class PDFProcessor:
             verbose=False
         )
     
-    def extract_layout(self, pdf_path: str) -> List[Layout]:
-        """Extract structured layout using MinerU"""
-        return self.mineru.process_pdf(pdf_path)
+    def process_pdf(self, pdf_path: str) -> Dict:
+        """Process PDF using MinerU pipeline"""
+        start_time = time.time()
+        
+        # Initialize MinerU components
+        local_image_dir = self.output_dir / "images"
+        local_md_dir = self.output_dir
+        image_dir = str(local_image_dir.name)
+        
+        os.makedirs(local_image_dir, exist_ok=True)
+        
+        image_writer = FileBasedDataWriter(str(local_image_dir))
+        md_writer = FileBasedDataWriter(str(local_md_dir))
+        
+        # Read PDF
+        reader = FileBasedDataReader("")
+        pdf_bytes = reader.read(pdf_path)
+        
+        # Create dataset and process
+        ds = PymuDocDataset(pdf_bytes)
+        
+        if ds.classify() == SupportedPdfParseMethod.OCR:
+            infer_result = ds.apply(doc_analyze, ocr=True)
+            pipe_result = infer_result.pipe_ocr_mode(image_writer)
+        else:
+            infer_result = ds.apply(doc_analyze, ocr=False)
+            pipe_result = infer_result.pipe_txt_mode(image_writer)
+        
+        # Get structured content
+        middle_json = pipe_result.get_middle_json()
+        tables = self._extract_tables(middle_json)
+        text_blocks = self._extract_text_blocks(middle_json)
+        
+        # Process text blocks with LLM
+        products = []
+        for block in text_blocks:
+            product = self._process_text_block(block)
+            if product:
+                product.tables = tables
+                products.append(product.to_dict())
+        
+        logger.info(f"Processed {len(products)} products in {time.time()-start_time:.2f}s")
+        return {"products": products, "tables": tables}
 
-    def process_tables(self, tables: List[Table]) -> List[Dict]:
-        """Convert MinerU tables to structured format"""
-        return [{
-            "page": table.page_number,
-            "cells": table.cells,
-            "header": table.headers,
-            "content": table.content
-        } for table in tables]
+    def _extract_tables(self, middle_json: Dict) -> List[Dict]:
+        """Extract tables from MinerU's middle JSON"""
+        tables = []
+        for page in middle_json.get('pages', []):
+            for table in page.get('tables', []):
+                tables.append({
+                    "page": page.get('page_number'),
+                    "cells": table.get('cells', []),
+                    "header": table.get('header', []),
+                    "content": table.get('content', [])
+                })
+        return tables
 
-    def generate_query_prompt(self, text: str) -> str:
-        """Create optimized extraction prompt"""
+    def _extract_text_blocks(self, middle_json: Dict) -> List[str]:
+        """Extract text blocks from MinerU's middle JSON"""
+        text_blocks = []
+        for page in middle_json.get('pages', []):
+            for block in page.get('blocks', []):
+                if block.get('type') == 'text':
+                    text_blocks.append(block.get('text', ''))
+        return text_blocks
+
+    def _process_text_block(self, text: str) -> Optional[ProductSpec]:
+        """Process text block with LLM"""
+        prompt = self._generate_query_prompt(text)
+        
+        try:
+            response = self.llm.create_chat_completion(
+                messages=[{"role": "user", "content": prompt}],
+                temperature=0.1,
+                max_tokens=512
+            )
+            return self._parse_response(response['choices'][0]['message']['content'])
+        except Exception as e:
+            logger.warning(f"Error processing text block: {e}")
+            return None
+
+    def _generate_query_prompt(self, text: str) -> str:
+        """Generate extraction prompt"""
         return f"""Extract product specifications from this text:
 {text}
 
@@ -70,8 +140,8 @@ Return JSON format:
     "attributes": {{ "key": "value" }}
 }}"""
 
-    def parse_response(self, response: str) -> Optional[ProductSpec]:
-        """Robust JSON parsing with fallbacks"""
+    def _parse_response(self, response: str) -> Optional[ProductSpec]:
+        """Parse LLM response"""
         try:
             json_start = response.find('{')
             json_end = response.rfind('}') + 1
@@ -86,33 +156,6 @@ Return JSON format:
             logger.warning(f"Parse error: {e}")
             return None
 
-    def process_pdf(self, pdf_path: str) -> Dict:
-        """Main processing pipeline"""
-        start_time = time.time()
-        
-        # Extract structured content
-        layout = self.extract_layout(pdf_path)
-        tables = self.process_tables(layout.tables)
-        
-        # Process text blocks
-        products = []
-        for block in layout.text_blocks:
-            prompt = self.generate_query_prompt(block.text)
-            
-            # Generate response with hardware optimization
-            response = self.llm.create_chat_completion(
-                messages=[{"role": "user", "content": prompt}],
-                temperature=0.1,
-                max_tokens=512
-            )
-            
-            if product := self.parse_response(response['choices'][0]['message']['content']):
-                product.tables = tables
-                products.append(product.to_dict())
-        
-        logger.info(f"Processed {len(products)} products in {time.time()-start_time:.2f}s")
-        return {"products": products, "tables": tables}
-
 def process_pdf_catalog(pdf_path: str):
     processor = PDFProcessor()
     try: