main.py

import os
import json
import time
import logging
from pathlib import Path
from typing import List, Dict, Optional
from dataclasses import dataclass
from fastapi.encoders import jsonable_encoder
from sentence_transformers import SentenceTransformer
from llama_cpp import Llama

# Fix: Dynamically adjust the module path if magic_pdf is in a non-standard location
try:
    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
except ModuleNotFoundError as e:
    logging.error(f"Failed to import magic_pdf modules: {e}")
    logging.info("Ensure that the magic_pdf package is installed and accessible in your Python environment.")
    raise e

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class ProductSpec:
    name: str
    description: Optional[str] = None
    price: Optional[float] = None
    attributes: Dict[str, str] = None
    tables: List[Dict] = None

    def to_dict(self):
        return jsonable_encoder(self)

class PDFProcessor:
    def __init__(self):
        self.emb_model = self._initialize_emb_model("all-MiniLM-L6-v2")
        self.llm = self._initialize_llm("deepseek-llm-7b-base.Q5_K_M.gguf")
        self.output_dir = Path("./output")
        self.output_dir.mkdir(exist_ok=True)

    def _initialize_emb_model(self, model_name):
        # model = SentenceTransformer("sentence-transformers/" + model_name)
        # model = SentenceTransformer(model_name)
        # model.save('models/'+ model_name)
        # Load model directly
        from transformers import AutoTokenizer, AutoModel

        tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
        model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
        return model

    def _initialize_llm(self, model_name):
        """Initialize LLM with automatic download if needed"""
        model_path = os.path.join("models/", model_name)
        if os.path.exists(model_path):
            return Llama(
                model_path=model_path,
                n_ctx=4096,
                n_gpu_layers=35 if os.getenv('USE_GPU') else 0,
                n_threads=os.cpu_count() - 1,
                verbose=False
            )
        else:
            return Llama.from_pretrained(
                repo_id="TheBloke/deepseek-llm-7B-base-GGUF",
                filename=model_name,
                n_ctx=4096,
                n_threads=os.cpu_count() - 1,
                n_gpu_layers=35 if os.getenv('USE_GPU') else 0,
                verbose=False
            )
    
    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)
        
        try:
            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}
        except Exception as e:
            logger.error(f"Error during PDF processing: {e}")
            raise RuntimeError("PDF processing failed.") from e
    
    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 _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}
Return JSON format:
{{
    "name": "product name",
    "description": "product description",
    "price": numeric_price,
    "attributes": {{ "key": "value" }}
}}"""
    
    def _parse_response(self, response: str) -> Optional[ProductSpec]:
        """Parse LLM response"""
        try:
            json_start = response.find('{')
            json_end = response.rfind('}') + 1
            data = json.loads(response[json_start:json_end])
            return ProductSpec(
                name=data.get('name', ''),
                description=data.get('description'),
                price=data.get('price'),
                attributes=data.get('attributes', {})
            )
        except (json.JSONDecodeError, KeyError) as e:
            logger.warning(f"Parse error: {e}")
            return None

def process_pdf_catalog(pdf_path: str):
    processor = PDFProcessor()
    try:
        result = processor.process_pdf(pdf_path)
        return result, "Processing completed successfully!"
    except Exception as e:
        logger.error(f"Processing failed: {e}")
        return {}, "Error processing PDF"