Update app.py

app.py

@@ -1,325 +1,169 @@
-# import gradio as gr
-# print("GRADIO VERSION:", gr.__version__)
-# import json
-# import os
-# import tempfile
-# from pathlib import Path
-
-# # NOTE: You must ensure that 'working_yolo_pipeline.py' exists 
-# # and defines the following items correctly:
-# from working_yolo_pipeline import run_document_pipeline, DEFAULT_LAYOUTLMV3_MODEL_PATH, WEIGHTS_PATH
-# # Since I don't have this file, I am assuming the imports are correct.
-
-# # Define placeholders for assumed constants if the pipeline file isn't present
-# # You should replace these with your actual definitions if they are missing
-# try:
-#     from working_yolo_pipeline import run_document_pipeline, DEFAULT_LAYOUTLMV3_MODEL_PATH, WEIGHTS_PATH
-# except ImportError:
-#     print("Warning: 'working_yolo_pipeline.py' not found. Using dummy paths.")
-#     def run_document_pipeline(*args):
-#         return {"error": "Placeholder pipeline function called."}
-#     DEFAULT_LAYOUTLMV3_MODEL_PATH = "./models/layoutlmv3_model"
-#     WEIGHTS_PATH = "./weights/yolo_weights.pt"
-
-
-# def process_pdf(pdf_file, layoutlmv3_model_path=None):
-#     """
-#     Wrapper function for Gradio interface.
-
-#     Args:
-#         pdf_file: Gradio UploadButton file object
-#         layoutlmv3_model_path: Optional custom model path
-
-#     Returns:
-#         Tuple of (JSON string, download file path)
-#     """
-#     if pdf_file is None:
-#         return "❌ Error: No PDF file uploaded.", None
-
-#     # Use default model path if not provided
-#     if not layoutlmv3_model_path:
-#         layoutlmv3_model_path = DEFAULT_LAYOUTLMV3_MODEL_PATH
-
-#     # Verify model and weights exist
-#     if not os.path.exists(layoutlmv3_model_path):
-#         return f"❌ Error: LayoutLMv3 model not found at {layoutlmv3_model_path}", None
-
-#     if not os.path.exists(WEIGHTS_PATH):
-#         return f"❌ Error: YOLO weights not found at {WEIGHTS_PATH}", None
-
-#     try:
-#         # Get the uploaded PDF path
-#         pdf_path = pdf_file.name
-
-#         # Run the pipeline
-#         result = run_document_pipeline(pdf_path, layoutlmv3_model_path, 'label_studio_import.json')
-
-#         if result is None:
-#             return "❌ Error: Pipeline failed to process the PDF. Check console for details.", None
-
-#         # Create a temporary file for download
-#         output_filename = f"{Path(pdf_path).stem}_analysis.json"
-#         temp_output = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', prefix='analysis_')
-
-#         # Dump results to the temporary file
-#         with open(temp_output.name, 'w', encoding='utf-8') as f:
-#             json.dump(result, f, indent=2, ensure_ascii=False)
-
-#         # Format JSON for display
-#         json_display = json.dumps(result, indent=2, ensure_ascii=False)
-
-#         return json_display, temp_output.name
-
-#     except Exception as e:
-#         return f"❌ Error during processing: {str(e)}", None
-
-
-# # Create Gradio interface
-# # FIX APPLIED: Removed 'theme=gr.themes.Soft()' which caused the TypeError
-# with gr.Blocks(title="Document Analysis Pipeline") as demo:
-#     gr.Markdown("""
-#     # 📄 Document Analysis Pipeline
-
-#     Upload a PDF document to extract structured data including questions, options, answers, passages, and embedded images.
-
-#     **Pipeline Steps:**
-#     1. 🔍 YOLO/OCR Preprocessing (word extraction + figure/equation detection)
-#     2. 🤖 LayoutLMv3 Inference (BIO tagging)
-#     3. 📊 Structured JSON Decoding
-#     4. 🖼️ Base64 Image Embedding
-#     """)
-
-#     with gr.Row():
-#         with gr.Column(scale=1):
-#             pdf_input = gr.File(
-#                 label="Upload PDF Document",
-#                 file_types=[".pdf"],
-#                 type="filepath"
-#             )
-
-#             model_path_input = gr.Textbox(
-#                 label="LayoutLMv3 Model Path (optional)",
-#                 placeholder=DEFAULT_LAYOUTLMV3_MODEL_PATH,
-#                 value=DEFAULT_LAYOUTLMV3_MODEL_PATH,
-#                 interactive=True
-#             )
-
-#             process_btn = gr.Button("🚀 Process Document", variant="primary", size="lg")
-
-#             gr.Markdown("""
-#             ### ℹ️ Notes:
-#             - Processing may take several minutes depending on PDF size
-#             - Figures and equations will be extracted and embedded as Base64
-#             - The output JSON includes structured questions, options, and answers
-#             """)
-
-#         with gr.Column(scale=2):
-#             json_output = gr.Code(
-#                 label="Structured JSON Output",
-#                 language="json",
-#                 lines=25
-#             )
-
-#             download_output = gr.File(
-#                 label="Download Full JSON",
-#                 interactive=False
-#             )
-
-#     # Status/Examples section
-#     with gr.Row():
-#         gr.Markdown("""
-#         ### 📋 Output Format
-#         The pipeline generates JSON with the following structure:
-#         - **Questions**: Extracted question text
-#         - **Options**: Multiple choice options (A, B, C, D, etc.)
-#         - **Answers**: Correct answer(s)
-#         - **Passages**: Associated reading passages
-#         - **Images**: Base64-encoded figures and equations (embedded with keys like `figure1`, `equation2`)
-#         """)
-
-#     # Connect the button to the processing function
-#     process_btn.click(
-#         fn=process_pdf,
-#         inputs=[pdf_input, model_path_input],
-#         outputs=[json_output, download_output],
-#         api_name="process_document"
-#     )
-
-#     # Example section (optional - add example PDFs if available)
-#     # gr.Examples(
-#     #     examples=[
-#     #         ["examples/sample1.pdf"],
-#     #         ["examples/sample2.pdf"],
-#     #     ],
-#     #     inputs=pdf_input,
-#     # )
-
-# # Launch the app
-# if __name__ == "__main__":
-#     demo.launch(
-#         server_name="0.0.0.0",
-#         server_port=7860,
-#         share=False,
-#         show_error=True
-#     )
-
-
-
-
-
 import gradio as gr
-print("GRADIO VERSION:", gr.__version__)
-import json
-import os
-import tempfile
-from pathlib import Path
-
-# ==============================
-# WRITE CUSTOM CSS FOR FONTS
-# ==============================
-
-# CUSTOM_CSS = """
-# @font-face {
-#     font-family: 'NotoSansMath';
-#     src: url('./NotoSansMath-Regular.ttf') format('truetype');
-#     font-weight: normal;
-#     font-style: normal;
-# }
-
-# html, body, * {
-#     font-family: 'NotoSansMath', sans-serif !important;
-# }
-# """
-
-# # Optionally write the CSS file if needed (not required for inline css)
-# if not os.path.exists("custom.css"):
-#     with open("custom.css", "w") as f:
-#         f.write(CUSTOM_CSS)
-# ==============================
-
-try:
-    from working_yolo_pipeline import run_document_pipeline, DEFAULT_LAYOUTLMV3_MODEL_PATH, WEIGHTS_PATH
-except ImportError:
-    print("Warning: 'working_yolo_pipeline.py' not found. Using dummy paths.")
-    def run_document_pipeline(*args):
-        return {"error": "Placeholder pipeline function called."}
-    DEFAULT_LAYOUTLMV3_MODEL_PATH = "./models/layoutlmv3_model"
-    WEIGHTS_PATH = "./weights/yolo_weights.pt"
-
-
-def process_pdf(pdf_file, layoutlmv3_model_path=None):
-    if pdf_file is None:
-        return "❌ Error: No PDF file uploaded.", None
-
-    if not layoutlmv3_model_path:
-        layoutlmv3_model_path = DEFAULT_LAYOUTLMV3_MODEL_PATH
-
-    if not os.path.exists(layoutlmv3_model_path):
-        return f"❌ Error: LayoutLMv3 model not found at {layoutlmv3_model_path}", None
-
-    if not os.path.exists(WEIGHTS_PATH):
-        return f"❌ Error: YOLO weights not found at {WEIGHTS_PATH}", None
-
-    try:
-        pdf_path = pdf_file.name
-
-        result = run_document_pipeline(pdf_path, layoutlmv3_model_path, 'label_studio_import.json')
-
-        if result is None:
-            return "❌ Error: Pipeline failed to process the PDF. Check console for details.", None
-
-        output_filename = f"{Path(pdf_path).stem}_analysis.json"
-        temp_output = tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json', prefix='analysis_')
-
-        with open(temp_output.name, 'w', encoding='utf-8') as f:
-            json.dump(result, f, indent=2, ensure_ascii=False)
-
-        json_display = json.dumps(result, indent=2, ensure_ascii=False)
-
-        return json_display, temp_output.name
-
-    except Exception as e:
-        return f"❌ Error during processing: {str(e)}", None
-
-
-with gr.Blocks(
-    title="Document Analysis Pipeline"
-) as demo:
-
-
-    gr.HTML()
-
-    gr.Markdown("""
-    # 📄 Document Analysis Pipeline
-
-    Upload a PDF document to extract structured data including questions, options, answers, passages, and embedded images.
-
-    **Pipeline Steps:**
-    1. 🔍 YOLO/OCR Preprocessing (word extraction + figure/equation detection)
-    2. 🤖 LayoutLMv3 Inference (BIO tagging)
-    3. 📊 Structured JSON Decoding
-    4. 🖼️ Base64 Image Embedding
-    """)
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            pdf_input = gr.File(
-                label="Upload PDF Document",
-                file_types=[".pdf"],
-                type="filepath"
-            )
-
-            model_path_input = gr.Textbox(
-                label="LayoutLMv3 Model Path (optional)",
-                placeholder=DEFAULT_LAYOUTLMV3_MODEL_PATH,
-                value=DEFAULT_LAYOUTLMV3_MODEL_PATH,
-                interactive=True
+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", ".", " ", ""]
             )
-
-            process_btn = gr.Button("🚀 Process Document", variant="primary", size="lg")
-
-            gr.Markdown("""
-            ### ℹ️ Notes:
-            - Processing may take several minutes depending on PDF size
-            - Figures and equations will be extracted and embedded as Base64
-            - The output JSON includes structured questions, options, and answers
-            """)
-
-        with gr.Column(scale=2):
-            json_output = gr.Code(
-                label="Structured JSON Output",
-                language="json",
-                lines=25
-            )
-
-            download_output = gr.File(
-                label="Download Full JSON",
-                interactive=False
+            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():
-        gr.Markdown("""
-        ### 📋 Output Format
-        The pipeline generates JSON with the following structure:
-        - **Questions**: Extracted question text
-        - **Options**: Multiple choice options
-        - **Answers**: Correct answer(s)
-        - **Passages**: Associated reading passages
-        - **Images**: Base64-encoded figures and equations
-        """)
+        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)
 
-    process_btn.click(
-        fn=process_pdf,
-        inputs=[pdf_input, model_path_input],
-        outputs=[json_output, download_output],
-        api_name="process_document"
-    )
+        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(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_error=True
-    )
+    demo.launch()