Update app.py

app.py

@@ -1,3 +1,4 @@
+
 # import gradio as gr
 # import fitz  # PyMuPDF
 # import torch
@@ -21,6 +22,9 @@
 # PROVIDERS = ["CPUExecutionProvider"]
 # print(f"⚡ Running on: {PROVIDERS}")
 
+
+
+
 # # ---------------------------------------------------------
 # # 1. OPTIMIZED EMBEDDINGS (BGE-SMALL)
 # # ---------------------------------------------------------
@@ -49,94 +53,6 @@
 #     def embed_query(self, text):
 #         return self._process_batch(["Represent this sentence for searching relevant passages: " + text])[0]
 
-# # ---------------------------------------------------------
-# # 2. OPTIMIZED LLM (Qwen 2.5 - 0.5B) - STRICT GRADING
-# # ---------------------------------------------------------
-# # class LLMEvaluator:
-# #     def __init__(self):
-# #         self.repo_id = "onnx-community/Qwen2.5-0.5B-Instruct" 
-# #         self.local_dir = "onnx_qwen_local"
-        
-# #         print(f"🔄 Preparing CPU LLM: {self.repo_id}...")
-        
-# #         if not os.path.exists(self.local_dir):
-# #             print(f"📥 Downloading FP16 model to {self.local_dir}...")
-# #             snapshot_download(
-# #                 repo_id=self.repo_id, 
-# #                 local_dir=self.local_dir,
-# #                 allow_patterns=["config.json", "generation_config.json", "tokenizer*", "special_tokens_map.json", "*.jinja", "onnx/model_fp16.onnx*"]
-# #             )
-# #             print("✅ Download complete.")
-
-# #         self.tokenizer = AutoTokenizer.from_pretrained(self.local_dir)
-        
-# #         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=False,
-# #             provider=PROVIDERS[0],
-# #             session_options=sess_options
-# #         )
-
-# #     def evaluate(self, context, question, student_answer, max_marks):
-# #         # OPTIMIZED PROMPT FOR SMALL MODELS (0.5B)
-# #         messages = [
-# #             {"role": "system", "content": "You are a strictest, literal academic grader in the whole. You ONLY grade based on the provided text. You DO NOT use outside knowledge."},
-# #             {"role": "user", "content": f"""
-# #             Task: Grade the student answer based ONLY on the Reference Text.
-            
-# #             REFERENCE TEXT:
-# #             {context}
-
-# #             QUESTION:
-# #             {question}
-
-# #             STUDENT ANSWER:
-# #             {student_answer}
-
-# #             -----------------------------
-# #             GRADING LOGIC:
-# #             1. READ the Reference Text and use that as the ground truth. What does it actually say about the Question?
-# #             2. COMPARE it to the Student Answer, do not forcefully agree with the answer by seeing things that are not there. You are to penalise irrelevant text and contradictions whenever you encounter them. 
-# #             3 START with 0 marks and IF the answers line up to the reference text in a meaningful way, then add marks porportionally. ONLY GIVE MARKS FOR CORRECT STATEMENT STRICTLY BASED ON THE REFERENCE TEXT AND NOTHING ELSE IN THIS WORLD.
-# #             4. IF the Student Answer claims things not found in the text , he is incorrect and HALLUCINATING. Do not give marks for that statment/phrase
-# #             5. IF the Student Answer contradicts the text (e.g., Text says "hide personality" but Student says "show personality"), Do not give marks for that statment/phrase
-            
-# #             VERDICT:
-# #             - If wrong: 0/{max_marks}
-# #             - If correct: {max_marks}/{max_marks}
-            
-# #             OUTPUT FORMAT:
-# #             Score: [X]/{max_marks}
-# #             Feedback: [Brief explanation citing the text]
-# #             """}
-# #         ]
-        
-# #         input_text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-# #         inputs = self.tokenizer(input_text, return_tensors="pt")
-        
-# #         with torch.no_grad():
-# #             outputs = self.model.generate(
-# #                 **inputs,
-# #                 max_new_tokens=100, 
-# #                 temperature=0.00,    # 0.0 = logic only, no creativity
-# #                 do_sample=False,
-# #                 repetition_penalty=1.2 
-# #             )
-        
-# #         input_length = inputs['input_ids'].shape[1]
-# #         response = self.tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
-# #         return response
-
-
-
-
-
 
 # # ---------------------------------------------------------
 # # 2. OPTIMIZED LLM (Qwen 2.5 - 0.5B) - STRICT GRADING
@@ -144,7 +60,7 @@
 # class LLMEvaluator:
 #     def __init__(self):
 #         # Qwen 2.5 0.5B is fast but needs "Few-Shot" examples to be strict.
-#         self.repo_id = "onnx-community/Qwen2.5-0.5B-Instruct" 
+#         self.repo_id = "onnx-community/Qwen2.5-1.5B-Instruct" 
 #         self.local_dir = "onnx_qwen_local"
         
 #         print(f"🔄 Preparing CPU LLM: {self.repo_id}...")
@@ -182,8 +98,8 @@
 #         Your job is to check if the Student Answer is FACTUALLY present in the Context.
 
 #         GRADING ALGORITHM:
-#         1. IF the Student Answer mentions things NOT in the Context -> PENALTY (-100%).
-#         2. IF the Student Answer interprets the text opposite to its meaning -> PENALTY (-100%).
+#         1. IF the Student Answer mentions things NOT in the Context -> PENALTY (-50% of the marks).
+#         2. IF the Student Answer interprets the text opposite to its meaning -> PENALTY (-100% of the marks).
 #         3. IF the Student Answer is generic fluff -> SCORE: 0.
 
 #         --- EXAMPLE 1 (HALLUCINATION) ---
@@ -248,110 +164,6 @@
 #         return response
 
 
-# # # ---------------------------------------------------------
-# # # 2. OPTIMIZED LLM (Qwen 2.5 - 0.5B) - STRICT GRADING
-# # # ---------------------------------------------------------
-# # class LLMEvaluator:
-# #     def __init__(self):
-# #         # Qwen 0.5B is great for speed, but needs VERY specific prompts to be strict.
-# #         self.repo_id = "onnx-community/Qwen2.5-0.5B-Instruct" 
-# #         self.local_dir = "onnx_qwen_local"
-        
-# #         print(f"🔄 Preparing CPU LLM: {self.repo_id}...")
-        
-# #         if not os.path.exists(self.local_dir):
-# #             print(f"📥 Downloading FP16 model to {self.local_dir}...")
-# #             snapshot_download(
-# #                 repo_id=self.repo_id, 
-# #                 local_dir=self.local_dir,
-# #                 allow_patterns=["config.json", "generation_config.json", "tokenizer*", "special_tokens_map.json", "*.jinja", "onnx/model_fp16.onnx*"]
-# #             )
-# #             print("✅ Download complete.")
-
-# #         self.tokenizer = AutoTokenizer.from_pretrained(self.local_dir)
-        
-# #         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=False,
-# #             provider=PROVIDERS[0],
-# #             session_options=sess_options
-# #         )
-
-# #     def evaluate(self, context, question, student_answer, max_marks):
-# #         # --- STRATEGY: FEW-SHOT PROMPTING & CHAIN OF THOUGHT ---
-# #         # Small models (0.5B) need examples to understand "Strictness".
-        
-# #         system_prompt = """You are a strict automated grader. You grade ONLY based on the provided Context.
-        
-# #         RULES:
-# #         1. If the Student Answer contains facts NOT found in the Context, Score is 0.
-# #         2. If the Student Answer contradicts the Context, Score is 0.
-# #         3. Do not use outside knowledge. If it's not in the text, it's wrong.
-        
-# #         --- EXAMPLE 1 (WRONG ANSWER) ---
-# #         Context: The sky is blue because of Rayleigh scattering.
-# #         Question: Why is the sky blue?
-# #         Student Answer: Because the ocean reflects into it.
-# #         Analysis: The context mentions Rayleigh scattering. The student mentioned ocean reflection. These do not match.
-# #         Score: 0/{max_marks}
-        
-# #         --- EXAMPLE 2 (CORRECT ANSWER) ---
-# #         Context: Mitochondria is the powerhouse of the cell.
-# #         Question: What is the mitochondria?
-# #         Student Answer: It is the powerhouse of the cell.
-# #         Analysis: The student answer matches the context text exactly.
-# #         Score: {max_marks}/{max_marks}
-# #         """
-
-# #         user_prompt = f"""
-# #         --- NOW GRADE THIS ---
-# #         Context: 
-# #         {context}
-
-# #         Question: 
-# #         {question}
-
-# #         Student Answer: 
-# #         {student_answer}
-
-# #         Task: 
-# #         1. Analyze if the specific keywords in Student Answer exist in Context.
-# #         2. Assign a Score.
-        
-# #         Output format:
-# #         Analysis: [Analysis here]
-# #         Score: [X]/{max_marks}
-# #         """
-
-# #         messages = [
-# #             {"role": "system", "content": system_prompt},
-# #             {"role": "user", "content": user_prompt}
-# #         ]
-        
-# #         input_text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-# #         inputs = self.tokenizer(input_text, return_tensors="pt")
-        
-# #         with torch.no_grad():
-# #             outputs = self.model.generate(
-# #                 **inputs,
-# #                 max_new_tokens=150, 
-# #                 temperature=0.1,    # Low temperature for facts
-# #                 top_p=0.1,          # Reduce creativity
-# #                 do_sample=True,
-# #                 repetition_penalty=1.1 
-# #             )
-        
-# #         input_length = inputs['input_ids'].shape[1]
-# #         response = self.tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
-# #         return response
-
-
 # # ---------------------------------------------------------
 # # 3. Main Application Logic
 # # ---------------------------------------------------------
@@ -363,23 +175,38 @@
 #         self.all_chunks = [] 
 #         self.total_chunks = 0
 
-#     def process_file(self, file_obj):
-#         if file_obj is None: return "No file uploaded."
+#     def process_content(self, file_obj, raw_text):
+#         # LOGIC: Check for exclusivity (Cannot have both file and text)
+#         has_file = file_obj is not None
+#         has_text = raw_text is not None and len(raw_text.strip()) > 0
+
+#         if has_file and has_text:
+#             return "❌ Error: Please provide EITHER a file OR paste text, not both at the same time."
+        
+#         if not has_file and not has_text:
+#             return "⚠️ No content provided. Please upload a file or paste text."
+
 #         try:
 #             text = ""
-#             if file_obj.name.endswith('.pdf'):
-#                 doc = fitz.open(file_obj.name)
-#                 for page in doc: text += page.get_text()
-#             elif file_obj.name.endswith('.txt'):
-#                 with open(file_obj.name, 'r', encoding='utf-8') as f: text = f.read()
+#             # Case 1: Process File
+#             if has_file:
+#                 if file_obj.name.endswith('.pdf'):
+#                     doc = fitz.open(file_obj.name)
+#                     for page in doc: text += page.get_text()
+#                 elif file_obj.name.endswith('.txt'):
+#                     with open(file_obj.name, 'r', encoding='utf-8') as f: text = f.read()
+#                 else:
+#                     return "❌ Error: Only .pdf and .txt supported."
+            
+#             # Case 2: Process Raw Text
 #             else:
-#                 return "❌ Error: Only .pdf and .txt supported."
+#                 text = raw_text
 
 #             text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
 #             self.all_chunks = text_splitter.split_text(text)
 #             self.total_chunks = len(self.all_chunks)
             
-#             if not self.all_chunks: return "File empty."
+#             if not self.all_chunks: return "Content empty."
 
 #             metadatas = [{"id": i} for i in range(self.total_chunks)]
 #             self.vector_store = FAISS.from_texts(self.all_chunks, self.embeddings, metadatas=metadatas)
@@ -389,7 +216,7 @@
 #             return f"Error: {str(e)}"
 
 #     def process_query(self, question, student_answer, max_marks):
-#         if not self.vector_store: return "⚠️ Please upload a file first.", ""
+#         if not self.vector_store: return "⚠️ Please upload a file or paste text first.", ""
 #         if not question: return "⚠️ Enter a question.", ""
 
 #         results = self.vector_store.similarity_search_with_score(question, k=1)
@@ -420,7 +247,11 @@
     
 #     with gr.Row():
 #         with gr.Column(scale=1):
-#             pdf_input = gr.File(label="1. Upload Chapter")
+#             gr.Markdown("### Source Input (Choose One)")
+#             pdf_input = gr.File(label="Option A: Upload Chapter (PDF/TXT)")
+#             gr.Markdown("**OR**")
+#             text_input = gr.Textbox(label="Option B: Paste Context", placeholder="Paste text here if you don't have a file...", lines=5)
+            
 #             upload_btn = gr.Button("Index Content", variant="primary")
 #             status_msg = gr.Textbox(label="Status", interactive=False)
 
@@ -436,7 +267,8 @@
 #                 evidence_box = gr.Markdown(label="Context Used")
 #                 grade_box = gr.Markdown(label="Grading Result")
 
-#     upload_btn.click(system.process_file, inputs=[pdf_input], outputs=[status_msg])
+#     # Pass both inputs to the process_content function
+#     upload_btn.click(system.process_content, inputs=[pdf_input, text_input], outputs=[status_msg])
 #     run_btn.click(system.process_query, inputs=[q_input, a_input, max_marks], outputs=[evidence_box, grade_box])
 
 # if __name__ == "__main__":
@@ -453,11 +285,19 @@
 
 
 
+
+
+
+
+
+
+
+
 import gradio as gr
 import fitz  # PyMuPDF
 import torch
 import os
-import onnxruntime as ort
+import numpy as np
 
 # --- IMPORT SESSION OPTIONS ---
 from onnxruntime import SessionOptions, GraphOptimizationLevel
@@ -466,10 +306,11 @@ from onnxruntime import SessionOptions, GraphOptimizationLevel
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
 from langchain_core.embeddings import Embeddings
+from langchain_core.documents import Document
 
 # --- ONNX & MODEL IMPORTS ---
 from transformers import AutoTokenizer
-from optimum.onnxruntime import ORTModelForFeatureExtraction, ORTModelForCausalLM
+from optimum.onnxruntime import ORTModelForFeatureExtraction, ORTModelForCausalLM, ORTModelForSequenceClassification
 from huggingface_hub import snapshot_download
 
 # Force CPU Provider
@@ -504,94 +345,6 @@ class OnnxBgeEmbeddings(Embeddings):
     def embed_query(self, text):
         return self._process_batch(["Represent this sentence for searching relevant passages: " + text])[0]
 
-# ---------------------------------------------------------
-# 2. OPTIMIZED LLM (Qwen 2.5 - 0.5B) - STRICT GRADING
-# ---------------------------------------------------------
-# class LLMEvaluator:
-#     def __init__(self):
-#         self.repo_id = "onnx-community/Qwen2.5-0.5B-Instruct" 
-#         self.local_dir = "onnx_qwen_local"
-        
-#         print(f"🔄 Preparing CPU LLM: {self.repo_id}...")
-        
-#         if not os.path.exists(self.local_dir):
-#             print(f"📥 Downloading FP16 model to {self.local_dir}...")
-#             snapshot_download(
-#                 repo_id=self.repo_id, 
-#                 local_dir=self.local_dir,
-#                 allow_patterns=["config.json", "generation_config.json", "tokenizer*", "special_tokens_map.json", "*.jinja", "onnx/model_fp16.onnx*"]
-#             )
-#             print("✅ Download complete.")
-
-#         self.tokenizer = AutoTokenizer.from_pretrained(self.local_dir)
-        
-#         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=False,
-#             provider=PROVIDERS[0],
-#             session_options=sess_options
-#         )
-
-#     def evaluate(self, context, question, student_answer, max_marks):
-#         # OPTIMIZED PROMPT FOR SMALL MODELS (0.5B)
-#         messages = [
-#             {"role": "system", "content": "You are a strictest, literal academic grader in the whole. You ONLY grade based on the provided text. You DO NOT use outside knowledge."},
-#             {"role": "user", "content": f"""
-#             Task: Grade the student answer based ONLY on the Reference Text.
-            
-#             REFERENCE TEXT:
-#             {context}
-
-#             QUESTION:
-#             {question}
-
-#             STUDENT ANSWER:
-#             {student_answer}
-
-#             -----------------------------
-#             GRADING LOGIC:
-#             1. READ the Reference Text and use that as the ground truth. What does it actually say about the Question?
-#             2. COMPARE it to the Student Answer, do not forcefully agree with the answer by seeing things that are not there. You are to penalise irrelevant text and contradictions whenever you encounter them. 
-#             3 START with 0 marks and IF the answers line up to the reference text in a meaningful way, then add marks porportionally. ONLY GIVE MARKS FOR CORRECT STATEMENT STRICTLY BASED ON THE REFERENCE TEXT AND NOTHING ELSE IN THIS WORLD.
-#             4. IF the Student Answer claims things not found in the text , he is incorrect and HALLUCINATING. Do not give marks for that statment/phrase
-#             5. IF the Student Answer contradicts the text (e.g., Text says "hide personality" but Student says "show personality"), Do not give marks for that statment/phrase
-            
-#             VERDICT:
-#             - If wrong: 0/{max_marks}
-#             - If correct: {max_marks}/{max_marks}
-            
-#             OUTPUT FORMAT:
-#             Score: [X]/{max_marks}
-#             Feedback: [Brief explanation citing the text]
-#             """}
-#         ]
-        
-#         input_text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-#         inputs = self.tokenizer(input_text, return_tensors="pt")
-        
-#         with torch.no_grad():
-#             outputs = self.model.generate(
-#                 **inputs,
-#                 max_new_tokens=100, 
-#                 temperature=0.00,    # 0.0 = logic only, no creativity
-#                 do_sample=False,
-#                 repetition_penalty=1.2 
-#             )
-        
-#         input_length = inputs['input_ids'].shape[1]
-#         response = self.tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
-#         return response
-
-
-
-
-
 
 # ---------------------------------------------------------
 # 2. OPTIMIZED LLM (Qwen 2.5 - 0.5B) - STRICT GRADING
@@ -630,9 +383,6 @@ class LLMEvaluator:
 
     def evaluate(self, context, question, student_answer, max_marks):
         # --- IMPROVED PROMPT STRATEGY ---
-        # 1. Role: We set the persona to a "Strict Logical Validator" not a "Teacher".
-        # 2. Few-Shot: We give examples of HALLUCINATIONS getting 0 marks.
-        
         system_prompt = f"""You are a strict Logic Validator. You are NOT a helpful assistant. 
         Your job is to check if the Student Answer is FACTUALLY present in the Context.
 
@@ -703,123 +453,70 @@ class LLMEvaluator:
         return response
 
 
-# # ---------------------------------------------------------
-# # 2. OPTIMIZED LLM (Qwen 2.5 - 0.5B) - STRICT GRADING
-# # ---------------------------------------------------------
-# class LLMEvaluator:
-#     def __init__(self):
-#         # Qwen 0.5B is great for speed, but needs VERY specific prompts to be strict.
-#         self.repo_id = "onnx-community/Qwen2.5-0.5B-Instruct" 
-#         self.local_dir = "onnx_qwen_local"
-        
-#         print(f"🔄 Preparing CPU LLM: {self.repo_id}...")
-        
-#         if not os.path.exists(self.local_dir):
-#             print(f"📥 Downloading FP16 model to {self.local_dir}...")
-#             snapshot_download(
-#                 repo_id=self.repo_id, 
-#                 local_dir=self.local_dir,
-#                 allow_patterns=["config.json", "generation_config.json", "tokenizer*", "special_tokens_map.json", "*.jinja", "onnx/model_fp16.onnx*"]
-#             )
-#             print("✅ Download complete.")
-
-#         self.tokenizer = AutoTokenizer.from_pretrained(self.local_dir)
-        
-#         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=False,
-#             provider=PROVIDERS[0],
-#             session_options=sess_options
-#         )
+# ---------------------------------------------------------
+# 3. NEW: ONNX RERANKER (Cross-Encoder)
+# Uses existing 'optimum' & 'transformers' libs (No new deps)
+# ---------------------------------------------------------
+class OnnxReranker:
+    def __init__(self):
+        # TinyBERT is ~17MB and very fast on CPU
+        self.model_name = "Xenova/ms-marco-TinyBERT-L-2-v2"
+        print(f"🔄 Loading Reranker: {self.model_name}...")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        self.model = ORTModelForSequenceClassification.from_pretrained(
+            self.model_name,
+            export=False,
+            provider=PROVIDERS[0]
+        )
 
-#     def evaluate(self, context, question, student_answer, max_marks):
-#         # --- STRATEGY: FEW-SHOT PROMPTING & CHAIN OF THOUGHT ---
-#         # Small models (0.5B) need examples to understand "Strictness".
+    def rank(self, query, docs, top_k=3):
+        if not docs:
+            return []
         
-#         system_prompt = """You are a strict automated grader. You grade ONLY based on the provided Context.
-        
-#         RULES:
-#         1. If the Student Answer contains facts NOT found in the Context, Score is 0.
-#         2. If the Student Answer contradicts the Context, Score is 0.
-#         3. Do not use outside knowledge. If it's not in the text, it's wrong.
-        
-#         --- EXAMPLE 1 (WRONG ANSWER) ---
-#         Context: The sky is blue because of Rayleigh scattering.
-#         Question: Why is the sky blue?
-#         Student Answer: Because the ocean reflects into it.
-#         Analysis: The context mentions Rayleigh scattering. The student mentioned ocean reflection. These do not match.
-#         Score: 0/{max_marks}
+        # Prepare pairs: [query, doc_text]
+        pairs = [[query, doc.page_content] for doc in docs]
         
-#         --- EXAMPLE 2 (CORRECT ANSWER) ---
-#         Context: Mitochondria is the powerhouse of the cell.
-#         Question: What is the mitochondria?
-#         Student Answer: It is the powerhouse of the cell.
-#         Analysis: The student answer matches the context text exactly.
-#         Score: {max_marks}/{max_marks}
-#         """
-
-#         user_prompt = f"""
-#         --- NOW GRADE THIS ---
-#         Context: 
-#         {context}
-
-#         Question: 
-#         {question}
-
-#         Student Answer: 
-#         {student_answer}
-
-#         Task: 
-#         1. Analyze if the specific keywords in Student Answer exist in Context.
-#         2. Assign a Score.
-        
-#         Output format:
-#         Analysis: [Analysis here]
-#         Score: [X]/{max_marks}
-#         """
-
-#         messages = [
-#             {"role": "system", "content": system_prompt},
-#             {"role": "user", "content": user_prompt}
-#         ]
+        inputs = self.tokenizer(
+            pairs, 
+            padding=True, 
+            truncation=True, 
+            max_length=512, 
+            return_tensors="pt"
+        )
         
-#         input_text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-#         inputs = self.tokenizer(input_text, return_tensors="pt")
+        with torch.no_grad():
+            outputs = self.model(**inputs)
         
-#         with torch.no_grad():
-#             outputs = self.model.generate(
-#                 **inputs,
-#                 max_new_tokens=150, 
-#                 temperature=0.1,    # Low temperature for facts
-#                 top_p=0.1,          # Reduce creativity
-#                 do_sample=True,
-#                 repetition_penalty=1.1 
-#             )
+        # Get logits (Relevance scores)
+        # MS-Marco models typically output a single logit or [irrelevant, relevant]
+        logits = outputs.logits
+        if logits.shape[1] == 2:
+            scores = logits[:, 1] # Take the "relevant" class score
+        else:
+            scores = logits.flatten()
+            
+        # Sort docs by score (descending)
+        scores = scores.numpy().tolist()
+        doc_score_pairs = list(zip(docs, scores))
+        doc_score_pairs.sort(key=lambda x: x[1], reverse=True)
         
-#         input_length = inputs['input_ids'].shape[1]
-#         response = self.tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
-#         return response
+        # Return top K docs
+        return [doc for doc, score in doc_score_pairs[:top_k]]
 
 
 # ---------------------------------------------------------
-# 3. Main Application Logic
+# 4. Main Application Logic
 # ---------------------------------------------------------
 class VectorSystem:
     def __init__(self):
         self.vector_store = None
         self.embeddings = OnnxBgeEmbeddings()
         self.llm = LLMEvaluator()
+        self.reranker = OnnxReranker() # Initialize Reranker
         self.all_chunks = [] 
         self.total_chunks = 0
 
     def process_content(self, file_obj, raw_text):
-        # LOGIC: Check for exclusivity (Cannot have both file and text)
         has_file = file_obj is not None
         has_text = raw_text is not None and len(raw_text.strip()) > 0
 
@@ -831,7 +528,6 @@ class VectorSystem:
 
         try:
             text = ""
-            # Case 1: Process File
             if has_file:
                 if file_obj.name.endswith('.pdf'):
                     doc = fitz.open(file_obj.name)
@@ -840,19 +536,21 @@ class VectorSystem:
                     with open(file_obj.name, 'r', encoding='utf-8') as f: text = f.read()
                 else:
                     return "❌ Error: Only .pdf and .txt supported."
-            
-            # Case 2: Process Raw Text
             else:
                 text = raw_text
 
-            text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
-            self.all_chunks = text_splitter.split_text(text)
-            self.total_chunks = len(self.all_chunks)
+            # Smaller chunks for Reranking precision (500 chars)
+            text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
+            texts = text_splitter.split_text(text)
+            self.all_chunks = texts # Keep plain text list for reference
+            
+            # Create Document objects with metadata
+            docs = [Document(page_content=t, metadata={"id": i}) for i, t in enumerate(texts)]
+            self.total_chunks = len(docs)
             
-            if not self.all_chunks: return "Content empty."
+            if not docs: return "Content empty."
 
-            metadatas = [{"id": i} for i in range(self.total_chunks)]
-            self.vector_store = FAISS.from_texts(self.all_chunks, self.embeddings, metadatas=metadatas)
+            self.vector_store = FAISS.from_documents(docs, self.embeddings)
             
             return f"✅ Indexed {self.total_chunks} chunks."
         except Exception as e:
@@ -862,19 +560,20 @@ class VectorSystem:
         if not self.vector_store: return "⚠️ Please upload a file or paste text first.", ""
         if not question: return "⚠️ Enter a question.", ""
 
-        results = self.vector_store.similarity_search_with_score(question, k=1)
-        top_doc, score = results[0]
+        # Step A: Wide Net Retrieval (Get top 15 candidates)
+        # We fetch more than we need to ensure the answer is in the candidate pool
+        initial_docs = self.vector_store.similarity_search(question, k=15)
         
-        center_id = top_doc.metadata['id']
-        start_id = max(0, center_id - 1)
-        end_id = min(self.total_chunks - 1, center_id + 1)
+        # Step B: Rerank (Get top 3 best matches)
+        # The Cross-Encoder strictly judges relevance
+        top_docs = self.reranker.rank(question, initial_docs, top_k=3)
         
-        expanded_context = ""
-        for i in range(start_id, end_id + 1):
-            expanded_context += self.all_chunks[i] + "\n"
+        # Step C: Construct Context
+        # We merge the top 3 specific chunks
+        expanded_context = "\n\n---\n\n".join([d.page_content for d in top_docs])
 
-        evidence_display = f"### 📚 Expanded Context (Chunks {start_id} to {end_id}):\n"
-        evidence_display += f"> ... {expanded_context} ..."
+        evidence_display = f"### 📚 Optimized Context (Top {len(top_docs)} chunks after Reranking):\n"
+        evidence_display += f"> {expanded_context} ..."
         
         llm_feedback = "Please enter a student answer to grade."
         if student_answer:
@@ -886,7 +585,7 @@ system = VectorSystem()
 
 with gr.Blocks(title="EduGenius AI Grader") as demo:
     gr.Markdown("# ⚡ EduGenius: CPU Optimized RAG")
-    gr.Markdown("Powered by **Qwen-2.5-0.5B** and **BGE-Small** (ONNX Optimized)")
+    gr.Markdown("Powered by **Qwen-2.5-0.5B**, **BGE-Small** & **TinyBERT Reranker**")
     
     with gr.Row():
         with gr.Column(scale=1):