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)