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 and literally dont go off on tangents making charitable assumptions, interprett the context and then the student answer and if they are not talking about the same thing then just deduct his marks like a vindictive teacher."},
            {"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 butbe so extrememly strict and do not make charitable hallucinations dont assume anythig that is not there.
            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()