Update app.py

app.py

@@ -22,7 +22,7 @@ import pytesseract
 
 # Import our models
 from simple.rag import initialize_models, process_documents, create_embedding, chunk_text_hierarchical
-from simple.ner import extract_legal_entities
+from simple.ner import process_text as run_ner
 from simple.summarizer import summarize_legal_document
 
 # Configure logging
@@ -37,9 +37,12 @@ cleanup_task = None
 # Configuration
 MONGODB_URI = os.getenv("MONGODB_URI", "mongodb+srv://username:password@cluster.mongodb.net/")
 DATABASE_NAME = os.getenv("DATABASE_NAME", "legal_rag_system")
-HF_MODEL_ID = os.getenv("HF_MODEL_ID", "sentence-transformers/all-MiniLM-L6-v2")
+# Hardcode embedding model per request
+HF_MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
 GROQ_API_KEY = os.getenv("GROQ_API_KEY", None)
 SESSION_EXPIRE_HOURS = int(os.getenv("SESSION_EXPIRE_HOURS", "24"))
+# Optional HF token (if NER model is private)
+HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN") or os.getenv("HF_TOKEN")
 
 # Supported file types
 SUPPORTED_EXTENSIONS = {'.pdf', '.txt', '.docx', '.doc'}
@@ -89,10 +92,17 @@ async def startup_event():
         # Create indexes
         await create_indexes()
         
-        # Initialize ML models
-        logger.info("🤖 Loading ML models...")
+        # Initialize ML models (embeddings / retrieval backbone)
+        logger.info(f"🤖 Loading embedding model for RAG: {HF_MODEL_ID}")
         initialize_models(HF_MODEL_ID, GROQ_API_KEY)
-        logger.info("✅ Models loaded successfully")
+        logger.info(f"✅ Embedding model loaded: {HF_MODEL_ID}")
+
+        # Surface NER token presence (actual NER loads lazily in simple.ner)
+        if HUGGINGFACE_TOKEN:
+            os.environ["HUGGINGFACE_TOKEN"] = HUGGINGFACE_TOKEN
+            logger.info("🔐 HUGGINGFACE_TOKEN detected for NER model access")
+        else:
+            logger.info("ℹ️ No HUGGINGFACE_TOKEN provided (NER model assumed public)")
         
         # Start cleanup task
         cleanup_task = asyncio.create_task(periodic_cleanup())
@@ -250,13 +260,19 @@ async def process_document_pipeline(
             {"$set": {"status": "processing", "updated_at": datetime.utcnow()}}
         )
         
-        # Step 1: NER Processing
-        logger.info(f"🔍 Running NER for session {session_id}")
-        ner_results = extract_legal_entities(
+        # Step 1: NER Processing (spaCy pipeline from Hugging Face)
+        ner_model_id = "kn29/my-ner-model"
+        logger.info(f"🔍 Running NER for session {session_id} using model: {ner_model_id}")
+        ner_results = run_ner(
             text,
-            model_id="kn29/my-ner-model",
-            hf_token=os.getenv("HF_TOKEN")  # optional, if your model is private
+            model_id=ner_model_id
         )
+        if ner_results.get("error"):
+            logger.error(f"❌ NER failed for session {session_id}: {ner_results['error']}")
+        else:
+            logger.info(
+                f"✅ NER completed for session {session_id} • total_entities={ner_results.get('total_entities', 0)} • labels={len(ner_results.get('unique_labels', []))}"
+            )
                 
         # Store NER results
         await db.ner_results.insert_one({
@@ -267,7 +283,7 @@ async def process_document_pipeline(
         })
         
         # Step 2: Summarization
-        logger.info(f"📄 Running summarization for session {session_id}")
+        logger.info(f"📄 Running summarization for session {session_id} (Groq={'on' if GROQ_API_KEY else 'off'})")
         summary_results = summarize_legal_document(
             text, 
             max_sentences=5, 
@@ -283,7 +299,7 @@ async def process_document_pipeline(
         })
         
         # Step 3: Chunking and Embedding
-        logger.info(f"🧩 Creating chunks and embeddings for session {session_id}")
+        logger.info(f"🧩 Creating chunks and embeddings for session {session_id} using {HF_MODEL_ID}")
         chunks = chunk_text_hierarchical(text, filename)
         
         # Create embeddings and store chunks