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 helpful academic grader. 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)
        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)
        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 - INT8 QUANTIZED)
# ---------------------------------------------------------
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} (Int8 Quantized)...")
        print(f"📥 Downloading Int8 model 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_quantized.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_quantized.onnx",
            use_cache=True,
            use_io_binding=False
        )

    def evaluate(self, context, question, student_answer):
        # SIMPLIFIED PROMPT (Easier for 1B model to follow)
        messages = [
            {"role": "system", "content": "You are a strict teacher. Grade the answer based ONLY on the text provided."},
            {"role": "user", "content": f"""
            SOURCE TEXT:
            {context}

            QUESTION: {question}
            ANSWER: {student_answer}

            TASK:
            1. Does the answer match the Source Text? (Yes/No)
            2. Score (0-10)
            3. Explanation (1 sentence)
            """}
        ]
        
        input_text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = self.tokenizer(input_text, return_tensors="pt")

        # [CRITICAL FIX] GENERATION SETTINGS
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=150,       # Keep it short
                do_sample=False,          # Greedy decoding (No randomness)
                repetition_penalty=1.3,   # STRONG penalty to kill loops like "The The The"
                min_length=5              # Force it to start speaking
            )
        
        return self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)

# ---------------------------------------------------------
# 3. Main Application Logic
# ---------------------------------------------------------
class VectorSystem:
    def __init__(self):
        self.vector_store = None
        self.embeddings = OnnxBgeEmbeddings()
        self.llm = LLMEvaluator() 
        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. RAG RETRIEVAL
        results = self.vector_store.similarity_search_with_score(question, k=1)
        if not results: return "No relevant text found.", ""

        best_doc, score = results[0]
        chunk_id = best_doc.metadata['id']

        # 2. SMART CONTEXT MERGING (Fixes the "Double Text" bug)
        # We manually overlap checking is hard, so we just grab the raw text range if possible
        # Simple fix: Concatenate with a separator to break the loop
        
        prev_text = self.all_chunks[chunk_id - 1] if chunk_id > 0 else ""
        curr_text = self.all_chunks[chunk_id]
        next_text = self.all_chunks[chunk_id + 1] if chunk_id < len(self.all_chunks) - 1 else ""

        # [FIX] We use "..." to separate them clearly for the LLM
        full_context = f"PREVIOUS: {prev_text}\n...\nFOCUS: {curr_text}\n...\nNEXT: {next_text}"

        # 3. LLM EVALUATION
        llm_feedback = "Please enter a student answer to grade."
        if student_answer:
            llm_feedback = self.llm.evaluate(full_context, question, student_answer)

        # UI Display
        evidence_display = f"### 🎯 Best Match (Score: {score:.4f})\n"
        evidence_display += f"> **PREVIOUS:** ...{prev_text[-300:]}\n\n"
        evidence_display += f"> **CORE:** **{curr_text}**\n\n"
        evidence_display += f"> **NEXT:** {next_text[:300]}...\n"

        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")
    
    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 Used")
                grade_box = gr.Markdown(label="LLM 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])

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