Upload 2 files

app.py

@@ -41,7 +41,7 @@ demo = gr.Interface(
     examples=pdf_examples,
     title="Open Source PDF Catalog Parser",
     description="Efficient PDF catalog processing using fitz and OpenLLM",
-    article="Uses MinerU for layout analysis and Llama-CPP for structured extraction"
+    article="Uses PyMuPDF for layout analysis and Llama-CPP for structured extraction"
 )
 
 if __name__ == "__main__":

main.py

@@ -5,10 +5,10 @@ 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
-from fastapi.encoders import jsonable_encoder
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -29,34 +29,25 @@ class ProductSpec:
 class PDFProcessor:
     def __init__(self):
         self.emb_model = self._initialize_emb_model("all-MiniLM-L6-v2")
-        # self.llm = self._initialize_llm("llama-2-7b.Q2_K.gguf")
+        # 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:
-            from sentence_transformers import SentenceTransformer
+            # Use SentenceTransformer if available
             return SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-        except:
-            # Load model directly
+        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"""
-        # model_path = os.path.join("models/", model_name)
-        # if os.path.exists(model_path):
-        #     return Llama(
-        #         model_path=model_path,
-        #         n_ctx=1024,
-        #         n_gpu_layers=-1,
-        #         n_threads=os.cpu_count() - 1,
-        #         verbose=False
-        #     )
+        # 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,
@@ -67,43 +58,63 @@ class PDFProcessor:
         start_time = time.time()
 
         # Open PDF
-        doc = fitz.open(pdf_path)
+        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
+        # Extract text and tables from each page
         for page_num, page in enumerate(doc):
-            # Extract text blocks
-            text_blocks.extend(self._extract_text_blocks(page))
+            # 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
+            # Extract tables (if any)
             tables.extend(self._extract_tables(page, page_num))
 
-        # Process text blocks with LLM
+        # Process text blocks with LLM to extract product information
         products = []
-        for block in text_blocks:
+        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
-                products.append(product.to_dict())
+                # 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"""
+        """Extract text blocks from a PDF page using PyMuPDF's blocks method."""
         blocks = []
         for block in page.get_text("blocks"):
-            blocks.append(block[4])  # The text content is at index 4
+            # 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"""
+        """Extract tables from a PDF page using PyMuPDF's table extraction (if available)."""
         tables = []
         try:
             tab = page.find_tables()
-            if tab.tables:
-                for table_idx, table in enumerate(tab.tables):
+            if tab and hasattr(tab, 'tables') and tab.tables:
+                for table in tab.tables:
                     table_data = table.extract()
                     if table_data:
                         tables.append({
@@ -113,51 +124,50 @@ class PDFProcessor:
                             "content": table_data
                         })
         except Exception as e:
-            logger.warning(f"Error extracting tables from page {page_num}: {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 text block with LLM"""
+        """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 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" }}
-}}"""
+        """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 LLM response"""
+        """Parse the LLM's response to extract a product specification."""
         try:
             json_start = response.find('{')
             json_end = response.rfind('}') + 1
-            data = json.loads(response[json_start:json_end])
+            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) as e:
-            logger.warning(f"Parse error: {e}")
+        except (json.JSONDecodeError, KeyError, ValueError) as e:
+            logger.warning(f"Parse error: {e} in response: {response}")
             return None
 
 
@@ -169,3 +179,10 @@ def process_pdf_catalog(pdf_path: str):
     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)