Update app.py

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, HTTPException
 import pytesseract
 import cv2
 import os
@@ -18,29 +18,30 @@ import cachetools
 import hashlib
 from vllm import LLM
 
-app = LLMAPI()
+app = FastAPI()
 
 # Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(name)s - %(message)s')
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 # Set Tesseract path
-pytesseract.pytesseract_cmd = "/usr/bin/tesseract"
+pytesseract.pytesseract.tesseract_cmd = "/usr/bin/tesseract"
 
 # Initialize BitNet model for CPU-only
 try:
     llm = LLM(
         model="username/bitnet-finetuned-invoice",  # Replace with your fine-tuned BitNet model
-        device="cuda",  # Explicitly set to CPU
-        enforce_eager=True,  # Disable CUDA graph compilation for CPU
+        device="cpu",
+        enforce_eager=True,  # Disable CUDA graph compilation
+        max_model_len=2048,  # Adjust based on memory (16GB RAM)
     )
 except Exception as e:
     logger.error(f"Failed to load BitNet model: {str(e)}")
-    raise RuntimeError("Failed to initialize BitNet model")
+    raise HTTPException(status_code=500, detail="BitNet model initialization failed")
 
-# In-memory caches (1GB TTL)
-raw_text_cache = cachetools.TTLTTL(maxsize=100, ttl=3600)
-structured_data_cache = cachetools.TTLTTL(maxsize=100, ttl=3600)
+# In-memory caches (1-hour TTL)
+raw_text_cache = cachetools.TTLCache(maxsize=100, ttl=3600)
+structured_data_cache = cachetools.TTLCache(maxsize=100, ttl=3600)
 
 def log_memory_usage():
     """Log current memory usage."""
@@ -86,7 +87,7 @@ async def process_pdf_page(img, page_idx):
         logger.info(f"Completed OCR for PDF page {page_idx}, took {time.time() - start_time:.2f} seconds, {log_memory_usage()}")
         return page_text + "\n"
     except Exception as e:
-        logger.error(f"OCR failed for PDF page {page_idx}: {str(e)}, {log_memory_usage()}")
+        logger.error(f"OCR failed for PDF page {idx}: {str(e)}, {log_memory_usage()}")
         return ""
 
 async def process_with_bitnet(filename: str, raw_text: str):
@@ -182,8 +183,8 @@ Output JSON:
   }}
 }}
 """
-        output = llm.generate([{"role": "user", "content": prompt}])
-        json_str = output[0].text
+        outputs = llm.generate(prompts=[prompt])
+        json_str = outputs[0].outputs[0].text
         json_start = json_str.find("{")
         json_end = json_str.rfind("}") + 1
         structured_data = json.loads(json_str[json_start:json_end])