app.py

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
import onnxruntime as ort

# Check available hardware accelerators
PROVIDERS = ort.get_available_providers()
print(f"⚡ Hardware Acceleration Providers: {PROVIDERS}")

# ---------------------------------------------------------
# 1. OPTIMIZED EMBEDDINGS (BGE-SMALL)
# ---------------------------------------------------------
class OnnxBgeEmbeddings(Embeddings):
    # CHANGE 1: Switched to 'bge-small' (3x faster than large, similar accuracy)
    def __init__(self, model_name="BAAI/bge-small-en-v1.5"):
        print(f"🔄 Loading Faster Embeddings: {model_name}...")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        
        self.model = ORTModelForFeatureExtraction.from_pretrained(
            model_name, 
            export=False,
            provider=PROVIDERS[0] # Auto-select best hardware (CUDA/CoreML)
        )

    def _process_batch(self, texts):
        inputs = self.tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt")
        
        # Move inputs to same device as model if needed (mostly handled by Optimum)
        device = self.model.device
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = self.model(**inputs)
        
        embeddings = outputs.last_hidden_state[:, 0]
        embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
        # Detach from graph before converting to numpy
        return embeddings.cpu().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. OPTIMIZED LLM (Qwen 2.5 - 0.5B)
# ---------------------------------------------------------
class LLMEvaluator:
    def __init__(self):
        # CHANGE 2: Switched to Qwen 2.5 0.5B (Half the size of Llama 1B, very smart)
        self.repo_id = "Xenova/Qwen2.5-0.5B-Instruct" 
        self.local_dir = "onnx_qwen_local"
        
        print(f"🔄 Preparing Ultra-Fast LLM: {self.repo_id}...")
        
        if not os.path.exists(self.local_dir):
            print(f"📥 Downloading Model to {self.local_dir}...")
            # Note: Xenova repos usually have the ONNX ready, no complex wildcard needed
            snapshot_download(repo_id=self.repo_id, local_dir=self.local_dir)
            print("✅ Download complete.")

        self.tokenizer = AutoTokenizer.from_pretrained(self.local_dir)
        
        # CHANGE 3: Enabled IO Binding + Explicit Provider
        self.model = ORTModelForCausalLM.from_pretrained(
            self.local_dir,
            use_cache=True,
            use_io_binding=True, # CHANGE: Major speedup on GPU
            provider=PROVIDERS[0]
        )

    def evaluate(self, context, question, student_answer, max_marks):
        # Qwen uses ChatML format implicitly via tokenizer
        messages = [
            {"role": "system", "content": "You are a strict academic grader. Verify the student answer against the context. Be harsh. Do not halluncinate."},
            {"role": "user", "content": f"""
            CONTEXT: {context}
            QUESTION: {question}
            ANSWER: {student_answer}
            
            TASK: Grade out of {max_marks}.
            RULES:
            1. If wrong, 0 marks.
            2. Be strict.
            3. Format: 'Score: X/{max_marks} \n Feedback: ...'
            """}
        ]
        
        input_text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = self.tokenizer(input_text, return_tensors="pt")
        
        # Move inputs for IO Binding
        device = self.model.device
        inputs = {k: v.to(device) for k, v in inputs.items()}

        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=75,  # CHANGE 4: Reduced tokens (we only need a short score/feedback)
                temperature=0.1,
                do_sample=False
            )
        
        response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        return response

# ---------------------------------------------------------
# 3. Main Application Logic (Unchanged but uses new classes)
# ---------------------------------------------------------
class VectorSystem:
    def __init__(self):
        self.vector_store = None
        self.embeddings = OnnxBgeEmbeddings() # Uses new BGE-Small
        self.llm = LLMEvaluator() # Uses new Qwen 0.5B
        self.all_chunks = [] 
        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."

            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.", ""

        results = self.vector_store.similarity_search_with_score(question, k=1)
        top_doc, score = results[0]
        
        center_id = top_doc.metadata['id']
        start_id = max(0, center_id - 1)
        end_id = min(self.total_chunks - 1, center_id + 1)
        
        expanded_context = ""
        for i in range(start_id, end_id + 1):
            expanded_context += self.all_chunks[i] + "\n"

        evidence_display = f"### 📚 Expanded Context (Chunks {start_id} to {end_id}):\n"
        evidence_display += f"> ... {expanded_context} ..."
        
        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

system = VectorSystem()

with gr.Blocks(title="EduGenius AI Grader") as demo:
    gr.Markdown("# ⚡ EduGenius: Ultra-Fast RAG")
    gr.Markdown("Powered by **Qwen-2.5-0.5B** and **BGE-Small** (ONNX Optimized)")
    
    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()