Update app.py

app.py

@@ -2,6 +2,10 @@ import gradio as gr
 import fitz  # PyMuPDF
 import torch
 import os
+import onnxruntime as ort
+
+# --- IMPORT SESSION OPTIONS ---
+from onnxruntime import SessionOptions, GraphOptimizationLevel
 
 # --- LANGCHAIN & RAG IMPORTS ---
 from langchain_text_splitters import RecursiveCharacterTextSplitter
@@ -12,39 +16,31 @@ from langchain_core.embeddings import Embeddings
 from transformers import AutoTokenizer
 from optimum.onnxruntime import ORTModelForFeatureExtraction, ORTModelForCausalLM
 from huggingface_hub import snapshot_download
-import onnxruntime as ort
 
 # Check available hardware accelerators
 PROVIDERS = ort.get_available_providers()
 print(f"⚡ Hardware Acceleration Providers: {PROVIDERS}")
 
 # ---------------------------------------------------------
-# 1. OPTIMIZED EMBEDDINGS (BGE-SMALL)
+# 1. OPTIMIZED EMBEDDINGS (BGE-SMALL) - [KEEP THIS SAME]
 # ---------------------------------------------------------
 class OnnxBgeEmbeddings(Embeddings):
     def __init__(self):
-        # FIX 1: Use "Xenova/..." version which has pre-exported ONNX weights.
-        # The official "BAAI/..." repo is PyTorch-only and fails with export=False.
         model_name = "Xenova/bge-small-en-v1.5"
         print(f"🔄 Loading Embeddings: {model_name}...")
-        
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        
         self.model = ORTModelForFeatureExtraction.from_pretrained(
             model_name, 
-            export=False, # Now safe because Xenova repo has model.onnx
-            provider=PROVIDERS[0] 
+            export=False, 
+            provider=PROVIDERS[0]
         )
 
     def _process_batch(self, texts):
         inputs = self.tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt")
-        
         device = self.model.device
         inputs = {k: v.to(device) for k, v in inputs.items()}
-
         with torch.no_grad():
             outputs = self.model(**inputs)
-        
         embeddings = outputs.last_hidden_state[:, 0]
         embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
         return embeddings.cpu().numpy().tolist()
@@ -56,19 +52,18 @@ class OnnxBgeEmbeddings(Embeddings):
         return self._process_batch(["Represent this sentence for searching relevant passages: " + text])[0]
 
 # ---------------------------------------------------------
-# 2. OPTIMIZED LLM (Qwen 2.5 - 0.5B)
+# 2. OPTIMIZED LLM (Qwen 2.5 - 0.5B) - [FIXED]
 # ---------------------------------------------------------
 class LLMEvaluator:
     def __init__(self):
-        # FIX 2: Correct Repo ID for Qwen 2.5 ONNX
         self.repo_id = "onnx-community/Qwen2.5-0.5B-Instruct" 
         self.local_dir = "onnx_qwen_local"
         
         print(f"🔄 Preparing Ultra-Fast LLM: {self.repo_id}...")
         
+        # Download (same as before)
         if not os.path.exists(self.local_dir):
             print(f"📥 Downloading FP16 model + data to {self.local_dir}...")
-            # We download the 'onnx' subfolder specifically
             snapshot_download(
                 repo_id=self.repo_id, 
                 local_dir=self.local_dir,
@@ -78,14 +73,19 @@ class LLMEvaluator:
 
         self.tokenizer = AutoTokenizer.from_pretrained(self.local_dir)
         
-        # FIX 3: Point to the 'onnx' subfolder inside the downloaded directory
+        # --- CRITICAL FIX: DISABLE GRAPH OPTIMIZATIONS ---
+        # The model is already optimized. Re-optimizing it at runtime causes the crash.
+        sess_options = SessionOptions()
+        sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_DISABLE_ALL
+        
         self.model = ORTModelForCausalLM.from_pretrained(
             self.local_dir,
             subfolder="onnx", 
             file_name="model_fp16.onnx",
             use_cache=True,
             use_io_binding=True, 
-            provider=PROVIDERS[0]
+            provider=PROVIDERS[0],
+            session_options=sess_options  # <--- PASS THIS HERE
         )
 
     def evaluate(self, context, question, student_answer, max_marks):