Delete main.py

main.py

@@ -1,188 +0,0 @@
-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
-import fitz  # PyMuPDF
-from sentence_transformers import SentenceTransformer
-from llama_cpp import Llama
-
-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")
-        # Choose the appropriate model filename below; adjust if needed.
-        self.llm = self._initialize_llm("deepseek-llm-7b-base.Q2_K.gguf")
-        self.output_dir = Path("./output")
-        self.output_dir.mkdir(exist_ok=True)
-
-    def _initialize_emb_model(self, model_name):
-        try:
-            # Use SentenceTransformer if available
-            return SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-        except Exception as e:
-            logger.warning(f"SentenceTransformer failed: {e}. Falling back to transformers model.")
-            from transformers import AutoTokenizer, AutoModel
-            tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/" + model_name)
-            model = AutoModel.from_pretrained("sentence-transformers/" + model_name)
-            return model
-
-    def _initialize_llm(self, model_name):
-        """Initialize LLM with automatic download if needed"""
-        # Here we use from_pretrained so that if the model is missing locally it downloads it.
-        return Llama.from_pretrained(
-            repo_id="TheBloke/deepseek-llm-7B-base-GGUF",
-            filename=model_name,
-        )
-
-    def process_pdf(self, pdf_path: str) -> Dict:
-        """Process PDF using PyMuPDF"""
-        start_time = time.time()
-
-        # Open PDF
-        try:
-            doc = fitz.open(pdf_path)
-        except Exception as e:
-            logger.error(f"Failed to open PDF: {e}")
-            raise RuntimeError("Cannot open PDF file.") from e
-
-        text_blocks = []
-        tables = []
-
-        # Extract text and tables from each page
-        for page_num, page in enumerate(doc):
-            # Extract text blocks from page and filter out very short blocks (noise)
-            blocks = self._extract_text_blocks(page)
-            filtered = [block for block in blocks if len(block.strip()) >= 10]
-            logger.debug(f"Page {page_num + 1}: Extracted {len(blocks)} blocks, {len(filtered)} kept after filtering.")
-            text_blocks.extend(filtered)
-
-            # Extract tables (if any)
-            tables.extend(self._extract_tables(page, page_num))
-
-        # Process text blocks with LLM to extract product information
-        products = []
-        for idx, block in enumerate(text_blocks):
-            # Log the text block for debugging
-            logger.debug(f"Processing text block {idx}: {block[:100]}...")
-            product = self._process_text_block(block)
-            if product:
-                product.tables = tables
-                # Only add if at least one key (like name) is non-empty
-                if product.name or product.description or product.price or (
-                        product.attributes and len(product.attributes) > 0):
-                    products.append(product.to_dict())
-                else:
-                    logger.debug(f"LLM returned empty product for block {idx}.")
-            else:
-                logger.debug(f"No product extracted from block {idx}.")
-
-        logger.info(f"Processed {len(products)} products in {time.time() - start_time:.2f}s")
-        return {"products": products, "tables": tables}
-
-    def _extract_text_blocks(self, page) -> List[str]:
-        """Extract text blocks from a PDF page using PyMuPDF's blocks method."""
-        blocks = []
-        for block in page.get_text("blocks"):
-            # block[4] contains the text content
-            text = block[4].strip()
-            if text:
-                blocks.append(text)
-        return blocks
-
-    def _extract_tables(self, page, page_num: int) -> List[Dict]:
-        """Extract tables from a PDF page using PyMuPDF's table extraction (if available)."""
-        tables = []
-        try:
-            tab = page.find_tables()
-            if tab and hasattr(tab, 'tables') and tab.tables:
-                for table in tab.tables:
-                    table_data = table.extract()
-                    if table_data:
-                        tables.append({
-                            "page": page_num + 1,
-                            "cells": table_data,
-                            "header": table.header.names if table.header else [],
-                            "content": table_data
-                        })
-        except Exception as e:
-            logger.warning(f"Error extracting tables from page {page_num + 1}: {e}")
-        return tables
-
-    def _process_text_block(self, text: str) -> Optional[ProductSpec]:
-        """Process a text block with LLM to extract product specifications."""
-        prompt = self._generate_query_prompt(text)
-        logger.debug(f"Generated prompt: {prompt[:200]}...")
-        try:
-            response = self.llm.create_chat_completion(
-                messages=[{"role": "user", "content": prompt}],
-                temperature=0.1,
-                max_tokens=512
-            )
-            # Debug: log raw response
-            logger.debug(f"LLM raw response: {response}")
-            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 a prompt instructing the LLM to extract product information."""
-        return f"""Extract product specifications from the following text. If no product is found, return an empty JSON object with keys.\n\nText:\n{text}\n\nReturn JSON format exactly as:\n{{\n    \"name\": \"product name\",\n    \"description\": \"product description\",\n    \"price\": numeric_price,\n    \"attributes\": {{ \"key\": \"value\" }}\n}}"""
-
-    def _parse_response(self, response: str) -> Optional[ProductSpec]:
-        """Parse the LLM's response to extract a product specification."""
-        try:
-            json_start = response.find('{')
-            json_end = response.rfind('}') + 1
-            json_str = response[json_start:json_end].strip()
-            if not json_str:
-                raise ValueError("No JSON content found in response.")
-            data = json.loads(json_str)
-            # If the returned JSON is essentially empty, return None
-            if all(not data.get(key) for key in ['name', 'description', 'price', 'attributes']):
-                return None
-            return ProductSpec(
-                name=data.get('name', ''),
-                description=data.get('description'),
-                price=data.get('price'),
-                attributes=data.get('attributes', {})
-            )
-        except (json.JSONDecodeError, KeyError, ValueError) as e:
-            logger.warning(f"Parse error: {e} in response: {response}")
-            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"
-
-
-if __name__ == "__main__":
-    # Example usage: change this if you call process_pdf_catalog elsewhere
-    pdf_path = "path/to/your/pdf_file.pdf"
-    result, message = process_pdf_catalog(pdf_path)
-    print(result, message)