Update app.py

app.py

@@ -179,6 +179,228 @@
 
 
 
+# import gradio as gr
+# import fitz  # PyMuPDF
+# import torch
+# import os
+
+# # --- LANGCHAIN & RAG IMPORTS ---
+# from langchain_text_splitters import RecursiveCharacterTextSplitter
+# from langchain_community.vectorstores import FAISS
+# from langchain_core.embeddings import Embeddings
+
+# # --- ONNX & MODEL IMPORTS ---
+# from transformers import AutoTokenizer
+# from optimum.onnxruntime import ORTModelForFeatureExtraction, ORTModelForCausalLM
+# from huggingface_hub import snapshot_download
+
+# # ---------------------------------------------------------
+# # 1. Custom ONNX Embedding Class (BGE-Large)
+# # ---------------------------------------------------------
+# class OnnxBgeEmbeddings(Embeddings):
+#     def __init__(self, model_name="BAAI/bge-large-en-v1.5"):
+#         print(f"🔄 Loading Embeddings: {model_name}...")
+#         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+#         # Note: export=True will re-convert on every restart. 
+#         # For production, you'd want to save this permanently, but this works for now.
+#         self.model = ORTModelForFeatureExtraction.from_pretrained(model_name, export=True)
+
+#     def _process_batch(self, texts):
+#         inputs = self.tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt")
+#         with torch.no_grad():
+#             outputs = self.model(**inputs)
+#         # CLS pooling for BGE
+#         embeddings = outputs.last_hidden_state[:, 0]
+#         embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+#         return embeddings.numpy().tolist()
+
+#     def embed_documents(self, texts):
+#         return self._process_batch(texts)
+
+#     def embed_query(self, text):
+#         return self._process_batch(["Represent this sentence for searching relevant passages: " + text])[0]
+
+# # ---------------------------------------------------------
+# # 2. LLM Evaluator Class (Llama-3.2-1B ONNX)
+# # ---------------------------------------------------------
+
+# class LLMEvaluator:
+#     def __init__(self):
+#         self.repo_id = "onnx-community/Llama-3.2-1B-Instruct"
+#         self.local_dir = "onnx_llama_local"
+        
+#         print(f"🔄 Preparing LLM: {self.repo_id}...")
+        
+#         # [FIXED DOWNLOADER]
+#         print(f"📥 Downloading FP16 model + data to {self.local_dir}...")
+#         snapshot_download(
+#             repo_id=self.repo_id,
+#             local_dir=self.local_dir,
+#             local_dir_use_symlinks=False,
+#             allow_patterns=[
+#                 "config.json", 
+#                 "generation_config.json", 
+#                 "tokenizer*", 
+#                 "special_tokens_map.json", 
+#                 "*.jinja", 
+#                 "onnx/model_fp16.onnx*"  # WILDCARD '*' ensures we get .onnx AND .onnx_data
+#             ]
+#         )
+#         print("✅ Download complete.")
+
+#         self.tokenizer = AutoTokenizer.from_pretrained(self.local_dir)
+        
+#         # [CRITICAL FIX] 
+#         # Separating 'subfolder' and 'file_name' is required by Optimum
+#         self.model = ORTModelForCausalLM.from_pretrained(
+#             self.local_dir,
+#             subfolder="onnx",             # Point to the subfolder
+#             file_name="model_fp16.onnx",  # Just the filename
+#             use_cache=True,
+#             use_io_binding=False
+#         )
+
+#     def evaluate(self, context, question, student_answer):
+#         # Prompt Engineering for Llama 3
+#         messages = [
+#             {"role": "system", "content": "You are a strict academic. Grade the student answer based ONLY on the provided context."},
+#             {"role": "user", "content": f"""
+#             ### CONTEXT:
+#             {context}
+
+#             ### QUESTION:
+#             {question}
+
+#             ### STUDENT ANSWER:
+#             {student_answer}
+
+#             ### INSTRUCTIONS:
+#             1. Is the answer correct?
+#             2. Score out of 10.
+#             3. Explanation.
+#             """}
+#         ]
+        
+#         # Format input using the chat template
+#         input_text = self.tokenizer.apply_chat_template(
+#             messages, 
+#             tokenize=False, 
+#             add_generation_prompt=True
+#         )
+        
+#         inputs = self.tokenizer(input_text, return_tensors="pt")
+
+#         # Generate response
+#         with torch.no_grad():
+#             outputs = self.model.generate(
+#                 **inputs,
+#                 max_new_tokens=256,
+#                 temperature=0.3,
+#                 do_sample=True,
+#                 top_p=0.9
+#             )
+        
+#         # Decode response
+#         response = self.tokenizer.decode(
+#             outputs[0][inputs.input_ids.shape[1]:], 
+#             skip_special_tokens=True
+#         )
+#         return response
+        
+# # ---------------------------------------------------------
+# # 3. Main Application Logic
+# # ---------------------------------------------------------
+# class VectorSystem:
+#     def __init__(self):
+#         self.vector_store = None
+#         self.embeddings = OnnxBgeEmbeddings()
+#         self.llm = LLMEvaluator() # Initialize LLM
+#         self.all_chunks = []
+
+#     def process_file(self, file_obj):
+#         if file_obj is None: return "No file uploaded."
+#         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()
+#             else:
+#                 return "❌ Error: Only .pdf and .txt supported."
+
+#             text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=150)
+#             self.all_chunks = text_splitter.split_text(text)
+            
+#             if not self.all_chunks: return "File empty."
+
+#             metadatas = [{"id": i} for i in range(len(self.all_chunks))]
+#             self.vector_store = FAISS.from_texts(self.all_chunks, self.embeddings, metadatas=metadatas)
+#             return f"✅ Indexed {len(self.all_chunks)} chunks."
+#         except Exception as e:
+#             return f"Error: {str(e)}"
+
+#     def process_query(self, question, student_answer):
+#         if not self.vector_store: return "⚠️ Please upload a file first.", ""
+#         if not question: return "⚠️ Enter a question.", ""
+
+#         # 1. Retrieve
+#         results = self.vector_store.similarity_search_with_score(question, k=3)
+        
+#         # Prepare context for LLM
+#         context_text = "\n\n".join([doc.page_content for doc, _ in results])
+        
+#         # Prepare Evidence Output for UI
+#         evidence_display = "### 📚 Retrieved Context:\n"
+#         for i, (doc, score) in enumerate(results):
+#             evidence_display += f"**Chunk {i+1}** (Score: {score:.4f}):\n> {doc.page_content}\n\n"
+
+#         # 2. Evaluate (if answer provided)
+#         llm_feedback = "Please enter a student answer to grade."
+#         if student_answer:
+#             llm_feedback = self.llm.evaluate(context_text, question, student_answer)
+
+#         return evidence_display, llm_feedback
+
+# # Initialize
+# system = VectorSystem()
+
+# # --- GRADIO UI ---
+# with gr.Blocks(title="EduGenius AI Grader") as demo:
+#     gr.Markdown("# 🧠 EduGenius: RAG + LLM Grading")
+#     gr.Markdown("Powered by **BGE-Large** (Retrieval) and **Llama-3.2-1B** (Evaluation) - All ONNX Optimized.")
+
+#     with gr.Row():
+#         with gr.Column(scale=1):
+#             pdf_input = gr.File(label="1. Upload Chapter (PDF/TXT)")
+#             upload_btn = gr.Button("Index Content", variant="primary")
+#             status_msg = gr.Textbox(label="System Status", interactive=False)
+
+#         with gr.Column(scale=2):
+#             q_input = gr.Textbox(label="2. Question")
+#             a_input = gr.Textbox(label="3. Student Answer")
+#             run_btn = gr.Button("Retrieve & Grade", variant="secondary")
+            
+#             with gr.Row():
+#                 evidence_box = gr.Markdown(label="Context")
+#                 grade_box = gr.Markdown(label="LLM Evaluation")
+
+#     upload_btn.click(system.process_file, inputs=[pdf_input], outputs=[status_msg])
+#     run_btn.click(system.process_query, inputs=[q_input, a_input], outputs=[evidence_box, grade_box])
+
+# if __name__ == "__main__":
+#     demo.launch()
+
+
+
+
+
+
+
+
+
+
+
 import gradio as gr
 import fitz  # PyMuPDF
 import torch
@@ -201,9 +423,11 @@ class OnnxBgeEmbeddings(Embeddings):
     def __init__(self, model_name="BAAI/bge-large-en-v1.5"):
         print(f"🔄 Loading Embeddings: {model_name}...")
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        # Note: export=True will re-convert on every restart. 
-        # For production, you'd want to save this permanently, but this works for now.
-        self.model = ORTModelForFeatureExtraction.from_pretrained(model_name, export=True)
+        
+        # OPTIMIZATION: Removed export=True. 
+        # Loading a pre-exported model or caching it is much faster.
+        # If you don't have the ONNX version, run export=True ONCE, then set to False.
+        self.model = ORTModelForFeatureExtraction.from_pretrained(model_name, export=False)
 
     def _process_batch(self, texts):
         inputs = self.tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt")
@@ -223,7 +447,6 @@ class OnnxBgeEmbeddings(Embeddings):
 # ---------------------------------------------------------
 # 2. LLM Evaluator Class (Llama-3.2-1B ONNX)
 # ---------------------------------------------------------
-
 class LLMEvaluator:
     def __init__(self):
         self.repo_id = "onnx-community/Llama-3.2-1B-Instruct"
@@ -231,39 +454,31 @@ class LLMEvaluator:
         
         print(f"🔄 Preparing LLM: {self.repo_id}...")
         
-        # [FIXED DOWNLOADER]
-        print(f"📥 Downloading FP16 model + data to {self.local_dir}...")
-        snapshot_download(
-            repo_id=self.repo_id,
-            local_dir=self.local_dir,
-            local_dir_use_symlinks=False,
-            allow_patterns=[
-                "config.json", 
-                "generation_config.json", 
-                "tokenizer*", 
-                "special_tokens_map.json", 
-                "*.jinja", 
-                "onnx/model_fp16.onnx*"  # WILDCARD '*' ensures we get .onnx AND .onnx_data
-            ]
-        )
-        print("✅ Download complete.")
+        # Download usually only needs to happen once
+        if not os.path.exists(self.local_dir):
+            print(f"📥 Downloading FP16 model + data to {self.local_dir}...")
+            snapshot_download(
+                repo_id=self.repo_id,
+                local_dir=self.local_dir,
+                local_dir_use_symlinks=False,
+                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)
         
-        # [CRITICAL FIX] 
-        # Separating 'subfolder' and 'file_name' is required by Optimum
         self.model = ORTModelForCausalLM.from_pretrained(
             self.local_dir,
-            subfolder="onnx",             # Point to the subfolder
-            file_name="model_fp16.onnx",  # Just the filename
+            subfolder="onnx",
+            file_name="model_fp16.onnx",
             use_cache=True,
             use_io_binding=False
         )
 
-    def evaluate(self, context, question, student_answer):
-        # Prompt Engineering for Llama 3
+    def evaluate(self, context, question, student_answer, max_marks):
+        # OPTIMIZATION: Strict Grading Prompt
         messages = [
-            {"role": "system", "content": "You are a strict academic. Grade the student answer based ONLY on the provided context."},
+            {"role": "system", "content": "You are a strict academic grader. You must grade accurately based on the context provided."},
             {"role": "user", "content": f"""
             ### CONTEXT:
             {context}
@@ -274,39 +489,32 @@ class LLMEvaluator:
             ### STUDENT ANSWER:
             {student_answer}
 
-            ### INSTRUCTIONS:
-            1. Is the answer correct?
-            2. Score out of 10.
-            3. Explanation.
+            ### GRADING INSTRUCTIONS:
+            1. The maximum score for this question is {max_marks}.
+            2. If the answer is completely wrong, give 0.
+            3. If the answer is correct but missing details, deduct marks proportionally.
+            4. DO NOT hallucinate a score higher than {max_marks}.
+            
+            ### OUTPUT FORMAT:
+            Score: [Your Score] / {max_marks}
+            Feedback: [One sentence explanation]
             """}
         ]
         
-        # Format input using the chat template
-        input_text = self.tokenizer.apply_chat_template(
-            messages, 
-            tokenize=False, 
-            add_generation_prompt=True
-        )
-        
+        input_text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         inputs = self.tokenizer(input_text, return_tensors="pt")
 
-        # Generate response
         with torch.no_grad():
             outputs = self.model.generate(
                 **inputs,
-                max_new_tokens=256,
-                temperature=0.3,
-                do_sample=True,
-                top_p=0.9
+                max_new_tokens=150, # Reduced to improve speed
+                temperature=0.1,    # Lower temperature for stricter, less creative grading
+                do_sample=False     # Greedy decoding is faster and more deterministic for grading
             )
         
-        # Decode response
-        response = self.tokenizer.decode(
-            outputs[0][inputs.input_ids.shape[1]:], 
-            skip_special_tokens=True
-        )
+        response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
         return response
-        
+
 # ---------------------------------------------------------
 # 3. Main Application Logic
 # ---------------------------------------------------------
@@ -314,8 +522,9 @@ class VectorSystem:
     def __init__(self):
         self.vector_store = None
         self.embeddings = OnnxBgeEmbeddings()
-        self.llm = LLMEvaluator() # Initialize LLM
-        self.all_chunks = []
+        self.llm = LLMEvaluator()
+        self.all_chunks = [] # Stores raw text
+        self.total_chunks = 0
 
     def process_file(self, file_obj):
         if file_obj is None: return "No file uploaded."
@@ -329,36 +538,52 @@ class VectorSystem:
             else:
                 return "❌ Error: Only .pdf and .txt supported."
 
-            text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=150)
+            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."
 
-            metadatas = [{"id": i} for i in range(len(self.all_chunks))]
+            # OPTIMIZATION: Store explicit IDs to allow neighbor retrieval
+            metadatas = [{"id": i} for i in range(self.total_chunks)]
             self.vector_store = FAISS.from_texts(self.all_chunks, self.embeddings, metadatas=metadatas)
-            return f"✅ Indexed {len(self.all_chunks)} chunks."
+            
+            return f"✅ Indexed {self.total_chunks} chunks."
         except Exception as e:
             return f"Error: {str(e)}"
 
-    def process_query(self, question, student_answer):
+    def process_query(self, question, student_answer, max_marks):
         if not self.vector_store: return "⚠️ Please upload a file first.", ""
         if not question: return "⚠️ Enter a question.", ""
 
-        # 1. Retrieve
-        results = self.vector_store.similarity_search_with_score(question, k=3)
+        # 1. Retrieve ONLY Top 1 Chunk
+        results = self.vector_store.similarity_search_with_score(question, k=1)
+        top_doc, score = results[0]
         
-        # Prepare context for LLM
-        context_text = "\n\n".join([doc.page_content for doc, _ in results])
+        # 2. Context Expansion (Preceding + Succeeding Chunks)
+        # Get the ID of the best match
+        center_id = top_doc.metadata['id']
         
-        # Prepare Evidence Output for UI
-        evidence_display = "### 📚 Retrieved Context:\n"
-        for i, (doc, score) in enumerate(results):
-            evidence_display += f"**Chunk {i+1}** (Score: {score:.4f}):\n> {doc.page_content}\n\n"
+        # Calculate indices (handle boundaries)
+        start_id = max(0, center_id - 1)
+        end_id = min(self.total_chunks - 1, center_id + 1)
+        
+        # Fetch the contiguous text block
+        expanded_context = ""
+        context_indices = []
+        
+        for i in range(start_id, end_id + 1):
+            expanded_context += self.all_chunks[i] + "\n"
+            context_indices.append(i)
 
-        # 2. Evaluate (if answer provided)
+        # UI Evidence Display
+        evidence_display = f"### 📚 Expanded Context (Chunks {start_id} to {end_id}):\n"
+        evidence_display += f"> ... {expanded_context} ..."
+        
+        # 3. Evaluate
         llm_feedback = "Please enter a student answer to grade."
         if student_answer:
-            llm_feedback = self.llm.evaluate(context_text, question, student_answer)
+            llm_feedback = self.llm.evaluate(expanded_context, question, student_answer, max_marks)
 
         return evidence_display, llm_feedback
 
@@ -367,27 +592,28 @@ system = VectorSystem()
 
 # --- GRADIO UI ---
 with gr.Blocks(title="EduGenius AI Grader") as demo:
-    gr.Markdown("# 🧠 EduGenius: RAG + LLM Grading")
-    gr.Markdown("Powered by **BGE-Large** (Retrieval) and **Llama-3.2-1B** (Evaluation) - All ONNX Optimized.")
-
+    gr.Markdown("# 🧠 EduGenius: Smart Context Grading")
+    
     with gr.Row():
         with gr.Column(scale=1):
-            pdf_input = gr.File(label="1. Upload Chapter (PDF/TXT)")
+            pdf_input = gr.File(label="1. Upload Chapter")
             upload_btn = gr.Button("Index Content", variant="primary")
-            status_msg = gr.Textbox(label="System Status", interactive=False)
+            status_msg = gr.Textbox(label="Status", interactive=False)
 
         with gr.Column(scale=2):
-            q_input = gr.Textbox(label="2. Question")
-            a_input = gr.Textbox(label="3. Student Answer")
+            with gr.Row():
+                q_input = gr.Textbox(label="Question", scale=2)
+                max_marks = gr.Slider(minimum=1, maximum=20, value=5, step=1, label="Max Marks")
+            
+            a_input = gr.TextArea(label="Student Answer")
             run_btn = gr.Button("Retrieve & Grade", variant="secondary")
             
             with gr.Row():
-                evidence_box = gr.Markdown(label="Context")
-                grade_box = gr.Markdown(label="LLM Evaluation")
+                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])
-    run_btn.click(system.process_query, inputs=[q_input, a_input], outputs=[evidence_box, grade_box])
+    run_btn.click(system.process_query, inputs=[q_input, a_input, max_marks], outputs=[evidence_box, grade_box])
 
 if __name__ == "__main__":
     demo.launch()
-