Update app.py

app.py

@@ -1,169 +1,361 @@
+# import gradio as gr
+# import fitz  # PyMuPDF
+# import torch
+# import numpy as np
+# from langchain_text_splitters import RecursiveCharacterTextSplitter
+# from langchain_community.vectorstores import FAISS
+# from langchain_core.embeddings import Embeddings
+
+# # --- NEW IMPORTS FOR ONNX ---
+# from transformers import AutoTokenizer
+# from optimum.onnxruntime import ORTModelForFeatureExtraction
+
+# # ---------------------------------------------------------
+# # Custom ONNX Embedding Class for BGE-Large
+# # ---------------------------------------------------------
+# class OnnxBgeEmbeddings(Embeddings):
+#     def __init__(self, model_name="BAAI/bge-large-en-v1.5", file_name="model.onnx"):
+#         print(f"🔄 Loading {model_name} with ONNX Runtime...")
+#         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        
+#         # This loads the model and exports it to ONNX format automatically if not already done
+#         self.model = ORTModelForFeatureExtraction.from_pretrained(
+#             model_name, 
+#             export=True 
+#         )
+#         self.model_name = model_name
+
+#     def _process_batch(self, texts):
+#         """Helper to tokenize and run inference via ONNX"""
+#         # Tokenize
+#         inputs = self.tokenizer(
+#             texts, 
+#             padding=True, 
+#             truncation=True, 
+#             max_length=512, 
+#             return_tensors="pt"
+#         )
+        
+#         # Run Inference (ONNX)
+#         with torch.no_grad():
+#             outputs = self.model(**inputs)
+        
+#         # BGE uses CLS pooling (first token), NOT mean pooling
+#         # outputs.last_hidden_state shape: [batch_size, seq_len, hidden_dim]
+#         embeddings = outputs.last_hidden_state[:, 0]
+        
+#         # Normalize embeddings (required for Cosine Similarity)
+#         embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
+        
+#         return embeddings.numpy().tolist()
+
+#     def embed_documents(self, texts):
+#         # BGE does NOT need instructions for documents
+#         return self._process_batch(texts)
+
+#     def embed_query(self, text):
+#         # BGE REQUIRES this specific instruction for queries to work best
+#         instruction = "Represent this sentence for searching relevant passages: "
+#         return self._process_batch([instruction + text])[0]
+
+# # ---------------------------------------------------------
+# # Main Application Logic
+# # ---------------------------------------------------------
+# class VectorSystem:
+#     def __init__(self):
+#         self.vector_store = None
+#         # SWITCHED to Custom ONNX Class
+#         self.embeddings = OnnxBgeEmbeddings(model_name="BAAI/bge-large-en-v1.5")
+#         self.all_chunks = []
+
+#     def process_file(self, file_obj):
+#         """Extracts text, preserves order, and builds the Vector Index"""
+#         if file_obj is None:
+#             return "No file uploaded."
+
+#         try:
+#             # 1. Extract Text
+#             text = ""
+#             file_path = file_obj.name
+            
+#             if file_path.lower().endswith('.pdf'):
+#                 doc = fitz.open(file_path)
+#                 for page in doc: text += page.get_text()
+#             elif file_path.lower().endswith('.txt'):
+#                 with open(file_path, 'r', encoding='utf-8') as f: text = f.read()
+#             else:
+#                 return "❌ Error: Only .pdf and .txt files are supported."
+
+#             # 2. Split Text
+#             # Adjusted chunk size slightly for the larger model context, but 800 is still good
+#             text_splitter = RecursiveCharacterTextSplitter(
+#                 chunk_size=800,
+#                 chunk_overlap=150,
+#                 separators=["\n\n", "\n", ".", " ", ""]
+#             )
+#             self.all_chunks = text_splitter.split_text(text)
+
+#             if not self.all_chunks:
+#                 return "Could not extract text. Is the file empty?"
+
+#             # 3. Build Vector Index
+#             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"✅ Success! Indexed {len(self.all_chunks)} chunks using BGE-Large (ONNX)."
+        
+#         except Exception as e:
+#             return f"Error processing file: {str(e)}"
+
+#     def retrieve_evidence(self, question, student_answer):
+#         if not self.vector_store:
+#             return "⚠️ Please upload and process a file first."
+
+#         if not question:
+#             return "⚠️ Please enter a Question."
+
+#         # BGE is very accurate, so we search for top 3
+#         results = self.vector_store.similarity_search_with_score(question, k=3)
+        
+#         output_text = "### 🔍 Expanded Context Analysis (Powered by BGE-Large ONNX):\n"
+        
+#         for i, (doc, score) in enumerate(results):
+#             chunk_id = doc.metadata['id']
+            
+#             prev_chunk = self.all_chunks[chunk_id - 1] if chunk_id > 0 else "(Start of Text)"
+#             next_chunk = self.all_chunks[chunk_id + 1] if chunk_id < len(self.all_chunks) - 1 else "(End of Text)"
+            
+#             # Note: FAISS returns L2 distance. Lower is better. 
+#             # With normalized vectors, L2 = 2 * (1 - CosineSimilarity).
+            
+#             output_text += f"\n#### 🎯 Match #{i+1} (Score: {score:.4f})\n"
+#             output_text += f"> **Preceding Context:**\n{prev_chunk}\n\n" 
+#             output_text += f"> **MATCH:**\n**{doc.page_content}**\n\n"
+#             output_text += f"> **Succeeding Context:**\n{next_chunk}\n"
+#             output_text += "---\n"
+
+#         return output_text
+
+# # Initialize System
+# system = VectorSystem()
+
+# # --- Gradio UI ---
+# with gr.Blocks(title="EduGenius Context Retriever") as demo:
+#     gr.Markdown("# 🎓 EduGenius: Smart Context Retriever")
+#     gr.Markdown("Upload a Chapter. Powered by **BGE-Large (ONNX Accelerated)** for superior accuracy.")
+
+#     with gr.Row():
+#         with gr.Column(scale=1):
+#             pdf_input = gr.File(label="1. Upload File (PDF or TXT)", file_types=[".pdf", ".txt"])
+#             upload_btn = gr.Button("Process File", variant="primary")
+#             upload_status = gr.Textbox(label="Status", interactive=False)
+
+#         with gr.Column(scale=2):
+#             question_input = gr.Textbox(label="2. Question", placeholder="e.g., What causes the chemical reaction?")
+#             answer_input = gr.Textbox(label="Student Answer (Optional)", placeholder="e.g., The heat causes it...")
+#             search_btn = gr.Button("Find Context + Neighbors", variant="secondary")
+            
+#             evidence_output = gr.Markdown(label="Relevant Text Chunks")
+
+#     upload_btn.click(fn=system.process_file, inputs=[pdf_input], outputs=[upload_status])
+#     search_btn.click(fn=system.retrieve_evidence, inputs=[question_input, answer_input], outputs=[evidence_output])
+
+# if __name__ == "__main__":
+#     demo.launch()
+
+
+
+
+
+
+
+
+
+
+
+
+
 import gradio as gr
 import fitz  # PyMuPDF
 import torch
-import numpy as np
+import os
+
+# --- LANGCHAIN & RAG IMPORTS ---
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
 from langchain_core.embeddings import Embeddings
 
-# --- NEW IMPORTS FOR ONNX ---
-from transformers import AutoTokenizer
-from optimum.onnxruntime import ORTModelForFeatureExtraction
+# --- ONNX & MODEL IMPORTS ---
+from transformers import AutoTokenizer, Pipeline
+from optimum.onnxruntime import ORTModelForFeatureExtraction, ORTModelForCausalLM
 
 # ---------------------------------------------------------
-# Custom ONNX Embedding Class for BGE-Large
+# 1. Custom ONNX Embedding Class (BGE-Large)
 # ---------------------------------------------------------
 class OnnxBgeEmbeddings(Embeddings):
     def __init__(self, model_name="BAAI/bge-large-en-v1.5", file_name="model.onnx"):
-        print(f"🔄 Loading {model_name} with ONNX Runtime...")
+        print(f"🔄 Loading Embeddings: {model_name}...")
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        
-        # This loads the model and exports it to ONNX format automatically if not already done
-        self.model = ORTModelForFeatureExtraction.from_pretrained(
-            model_name, 
-            export=True 
-        )
-        self.model_name = model_name
+        self.model = ORTModelForFeatureExtraction.from_pretrained(model_name, export=True)
 
     def _process_batch(self, texts):
-        """Helper to tokenize and run inference via ONNX"""
-        # Tokenize
-        inputs = self.tokenizer(
-            texts, 
-            padding=True, 
-            truncation=True, 
-            max_length=512, 
-            return_tensors="pt"
-        )
-        
-        # Run Inference (ONNX)
+        inputs = self.tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt")
         with torch.no_grad():
             outputs = self.model(**inputs)
-        
-        # BGE uses CLS pooling (first token), NOT mean pooling
-        # outputs.last_hidden_state shape: [batch_size, seq_len, hidden_dim]
+        # CLS pooling for BGE
         embeddings = outputs.last_hidden_state[:, 0]
-        
-        # Normalize embeddings (required for Cosine Similarity)
         embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
-        
         return embeddings.numpy().tolist()
 
     def embed_documents(self, texts):
-        # BGE does NOT need instructions for documents
         return self._process_batch(texts)
 
     def embed_query(self, text):
-        # BGE REQUIRES this specific instruction for queries to work best
-        instruction = "Represent this sentence for searching relevant passages: "
-        return self._process_batch([instruction + text])[0]
+        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):
+        # Using the ONNX Community version of Llama 3.2 1B (Optimized for CPU)
+        self.model_id = "onnx-community/Llama-3.2-1B-Instruct"
+        print(f"🔄 Loading LLM: {self.model_id}...")
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+        
+        # Load the ONNX model for text generation
+        self.model = ORTModelForCausalLM.from_pretrained(
+            self.model_id,
+            decoder_file_name="model.onnx", # Standard ONNX filename
+            use_cache=True,
+            use_io_binding=False # Safer for CPU spaces
+        )
+
+    def evaluate(self, context, question, student_answer):
+        # Prompt Engineering for Llama 3
+        messages = [
+            {"role": "system", "content": "You are a strict but helpful academic grader. You will be given a context, a question, and a student's answer. Your job is to grade the answer based ONLY on the provided context."},
+            {"role": "user", "content": f"""
+            ### CONTEXT:
+            {context}
+
+            ### QUESTION:
+            {question}
+
+            ### STUDENT ANSWER:
+            {student_answer}
+
+            ### INSTRUCTIONS:
+            1. Determine if the student answer is correct based on the context.
+            2. Give a score out of 10.
+            3. Provide a brief 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, # Low temp for factual grading
+                do_sample=True,
+                top_p=0.9
+            )
+        
+        # Decode and strip the prompt
+        response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+        return response
 
 # ---------------------------------------------------------
-# Main Application Logic
+# 3. Main Application Logic
 # ---------------------------------------------------------
 class VectorSystem:
     def __init__(self):
         self.vector_store = None
-        # SWITCHED to Custom ONNX Class
-        self.embeddings = OnnxBgeEmbeddings(model_name="BAAI/bge-large-en-v1.5")
+        self.embeddings = OnnxBgeEmbeddings()
+        self.llm = LLMEvaluator() # Initialize LLM
         self.all_chunks = []
 
     def process_file(self, file_obj):
-        """Extracts text, preserves order, and builds the Vector Index"""
-        if file_obj is None:
-            return "No file uploaded."
-
+        if file_obj is None: return "No file uploaded."
         try:
-            # 1. Extract Text
             text = ""
-            file_path = file_obj.name
-            
-            if file_path.lower().endswith('.pdf'):
-                doc = fitz.open(file_path)
+            if file_obj.name.endswith('.pdf'):
+                doc = fitz.open(file_obj.name)
                 for page in doc: text += page.get_text()
-            elif file_path.lower().endswith('.txt'):
-                with open(file_path, 'r', encoding='utf-8') as f: text = f.read()
+            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 files are supported."
-
-            # 2. Split Text
-            # Adjusted chunk size slightly for the larger model context, but 800 is still good
-            text_splitter = RecursiveCharacterTextSplitter(
-                chunk_size=800,
-                chunk_overlap=150,
-                separators=["\n\n", "\n", ".", " ", ""]
-            )
-            self.all_chunks = text_splitter.split_text(text)
+                return "❌ Error: Only .pdf and .txt supported."
 
-            if not self.all_chunks:
-                return "Could not extract text. Is the file empty?"
+            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."
 
-            # 3. Build Vector Index
             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"✅ Success! Indexed {len(self.all_chunks)} chunks using BGE-Large (ONNX)."
-        
+            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 processing file: {str(e)}"
+            return f"Error: {str(e)}"
 
-    def retrieve_evidence(self, question, student_answer):
-        if not self.vector_store:
-            return "⚠️ Please upload and process a file first."
+    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.", ""
 
-        if not question:
-            return "⚠️ Please enter a Question."
-
-        # BGE is very accurate, so we search for top 3
+        # 1. Retrieve
         results = self.vector_store.similarity_search_with_score(question, k=3)
         
-        output_text = "### 🔍 Expanded Context Analysis (Powered by BGE-Large ONNX):\n"
+        # 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):
-            chunk_id = doc.metadata['id']
-            
-            prev_chunk = self.all_chunks[chunk_id - 1] if chunk_id > 0 else "(Start of Text)"
-            next_chunk = self.all_chunks[chunk_id + 1] if chunk_id < len(self.all_chunks) - 1 else "(End of Text)"
-            
-            # Note: FAISS returns L2 distance. Lower is better. 
-            # With normalized vectors, L2 = 2 * (1 - CosineSimilarity).
-            
-            output_text += f"\n#### 🎯 Match #{i+1} (Score: {score:.4f})\n"
-            output_text += f"> **Preceding Context:**\n{prev_chunk}\n\n" 
-            output_text += f"> **MATCH:**\n**{doc.page_content}**\n\n"
-            output_text += f"> **Succeeding Context:**\n{next_chunk}\n"
-            output_text += "---\n"
+            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 output_text
+        return evidence_display, llm_feedback
 
-# Initialize System
+# Initialize
 system = VectorSystem()
 
-# --- Gradio UI ---
-with gr.Blocks(title="EduGenius Context Retriever") as demo:
-    gr.Markdown("# 🎓 EduGenius: Smart Context Retriever")
-    gr.Markdown("Upload a Chapter. Powered by **BGE-Large (ONNX Accelerated)** for superior accuracy.")
+# --- 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 File (PDF or TXT)", file_types=[".pdf", ".txt"])
-            upload_btn = gr.Button("Process File", variant="primary")
-            upload_status = gr.Textbox(label="Status", interactive=False)
+            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):
-            question_input = gr.Textbox(label="2. Question", placeholder="e.g., What causes the chemical reaction?")
-            answer_input = gr.Textbox(label="Student Answer (Optional)", placeholder="e.g., The heat causes it...")
-            search_btn = gr.Button("Find Context + Neighbors", variant="secondary")
+            q_input = gr.Textbox(label="2. Question")
+            a_input = gr.Textbox(label="3. Student Answer")
+            run_btn = gr.Button("Retrieve & Grade", variant="secondary")
             
-            evidence_output = gr.Markdown(label="Relevant Text Chunks")
+            with gr.Row():
+                evidence_box = gr.Markdown(label="Context")
+                grade_box = gr.Markdown(label="LLM Evaluation")
 
-    upload_btn.click(fn=system.process_file, inputs=[pdf_input], outputs=[upload_status])
-    search_btn.click(fn=system.retrieve_evidence, inputs=[question_input, answer_input], outputs=[evidence_output])
+    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()