app.py

# 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 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
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)
        
        # 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")
        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}...")
        
        # 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)
        
        self.model = ORTModelForCausalLM.from_pretrained(
            self.local_dir,
            subfolder="onnx",
            file_name="model_fp16.onnx",
            use_cache=True,
            use_io_binding=False
        )

    def evaluate(self, context, question, student_answer, max_marks):
        # OPTIMIZATION: Strict Grading Prompt
        messages = [
            {"role": "system", "content": "You are a strict academic grader. You must grade accurately based on the context provided."},
            {"role": "user", "content": f"""
            ### CONTEXT:
            {context}

            ### QUESTION:
            {question}

            ### STUDENT ANSWER:
            {student_answer}

            ### 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]
            """}
        ]
        
        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, # 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
            )
        
        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()
        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."
        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=100)
            self.all_chunks = text_splitter.split_text(text)
            self.total_chunks = len(self.all_chunks)
            
            if not self.all_chunks: return "File empty."

            # 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 {self.total_chunks} chunks."
        except Exception as e:
            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 question: return "⚠️ Enter a question.", ""

        # 1. Retrieve ONLY Top 1 Chunk
        results = self.vector_store.similarity_search_with_score(question, k=1)
        top_doc, score = results[0]
        
        # 2. Context Expansion (Preceding + Succeeding Chunks)
        # Get the ID of the best match
        center_id = top_doc.metadata['id']
        
        # 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)

        # 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(expanded_context, question, student_answer, max_marks)

        return evidence_display, llm_feedback

# Initialize
system = VectorSystem()

# --- GRADIO UI ---
with gr.Blocks(title="EduGenius AI Grader") as demo:
    gr.Markdown("# 🧠 EduGenius: Smart Context Grading")
    
    with gr.Row():
        with gr.Column(scale=1):
            pdf_input = gr.File(label="1. Upload Chapter")
            upload_btn = gr.Button("Index Content", variant="primary")
            status_msg = gr.Textbox(label="Status", interactive=False)

        with gr.Column(scale=2):
            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 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, max_marks], outputs=[evidence_box, grade_box])

if __name__ == "__main__":
    demo.launch()