Upload app.py

app.py

@@ -5,98 +5,69 @@ import fitz  # PyMuPDF
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.chains import RetrievalQA
-from langchain_community.llms import HuggingFacePipeline
-from langchain.prompts import PromptTemplate
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import torch
-from typing import List, Dict, Any
-import re
 import base64
 from PIL import Image
 import io
 
-class CurriculumAssistant:
-    def __init__(self):
+# --- Minimal PDF Search & Display App ---
+
+# 1. Preprocess PDFs and build vector DB
+class FastPDFSearch:
+    def __init__(self, slides_dir="Slides"):
+        self.pdf_pages = {}  # {filename: {page_num: text}}
+        self.pdf_files = {}  # {filename: path}
+        self.chunks = []
+        self.chunk_metadata = []
         self.vector_db = None
-        self.qa_chain = None
         self.embeddings = None
-        self.llm = None
-        self.curriculum_docs = []
-        self.pdf_pages = {}  # Store page-level information
-        self.pdf_files = {}  # Store PDF file objects for page rendering
-        
-    def load_llm(self):
-        """Load a better model for responses"""
-        try:
-            # Use a more capable model for better responses
-            model_name = "microsoft/DialoGPT-large"
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                torch_dtype=torch.float16,
-                device_map=None,  # Use CPU for Hugging Face Spaces
-                trust_remote_code=True
-            )
-            
-            pipe = pipeline(
-                "text-generation",
-                model=model,
-                tokenizer=tokenizer,
-                max_new_tokens=200,  # Increased for better responses
-                temperature=0.7,
-                top_p=0.95,
-                repetition_penalty=1.15,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id
-            )
-            
-            self.llm = HuggingFacePipeline(pipeline=pipe)
-            return True
-        except Exception as e:
-            print(f"Error loading model: {str(e)}")
-            return False
-    
-    def extract_text_from_pdf_with_pages(self, pdf_path: str) -> Dict[int, str]:
-        """Extract text from PDF file with page numbers"""
-        try:
-            doc = fitz.open(pdf_path)
+        self._process_pdfs(slides_dir)
+        self._build_vector_db()
+
+    def _process_pdfs(self, slides_dir):
+        slides_path = Path(slides_dir)
+        pdf_files = list(slides_path.glob("*.pdf"))
+        for pdf_file in pdf_files:
+            self.pdf_files[pdf_file.name] = str(pdf_file)
+            doc = fitz.open(str(pdf_file))
             pages = {}
             for page_num in range(len(doc)):
                 page = doc[page_num]
                 text = page.get_text()
-                if text.strip():  # Only store non-empty pages
+                if text.strip():
                     pages[page_num + 1] = text.strip()
+            self.pdf_pages[pdf_file.name] = pages
             doc.close()
-            return pages
-        except Exception as e:
-            print(f"Error extracting text from {pdf_path}: {str(e)}")
-            return {}
-    
-    def get_pdf_page_image(self, pdf_path: str, page_num: int) -> str:
-        """Get a specific page from PDF as base64 image"""
+            # Add each page as a chunk
+            for page_num, text in pages.items():
+                self.chunks.append(text)
+                self.chunk_metadata.append({
+                    "filename": pdf_file.name,
+                    "page_number": page_num
+                })
+
+    def _build_vector_db(self):
+        self.embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+        self.vector_db = Chroma.from_texts(
+            texts=self.chunks,
+            embedding=self.embeddings,
+            metadatas=self.chunk_metadata,
+            persist_directory="./chroma_db"
+        )
+
+    def get_pdf_page_image(self, pdf_path, page_num):
         try:
             doc = fitz.open(pdf_path)
             if page_num <= len(doc):
-                page = doc[page_num - 1]  # Convert to 0-based index
-                # Render page as image with higher quality
-                mat = fitz.Matrix(1.5, 1.5)  # Scale for better quality
+                page = doc[page_num - 1]
+                mat = fitz.Matrix(1.5, 1.5)
                 pix = page.get_pixmap(matrix=mat)
-                
-                # Convert to PIL Image for better handling
                 img_data = pix.tobytes("png")
                 img = Image.open(io.BytesIO(img_data))
-                
-                # Convert to RGB if needed
                 if img.mode != 'RGB':
                     img = img.convert('RGB')
-                
-                # Save to bytes
                 img_byte_arr = io.BytesIO()
                 img.save(img_byte_arr, format='PNG')
                 img_byte_arr = img_byte_arr.getvalue()
-                
-                # Convert to base64
                 img_base64 = base64.b64encode(img_byte_arr).decode()
                 doc.close()
                 return f"data:image/png;base64,{img_base64}"
@@ -105,299 +76,40 @@ class CurriculumAssistant:
         except Exception as e:
             print(f"Error rendering PDF page: {str(e)}")
             return None
-    
-    def process_curriculum(self, slides_dir: str):
-        """Process all PDF files in the slides directory"""
-        try:
-            slides_path = Path(slides_dir)
-            pdf_files = list(slides_path.glob("*.pdf"))
-            
-            if not pdf_files:
-                print("No PDF files found in the Slides directory!")
-                return False
-            
-            all_texts = []
-            all_chunks_with_metadata = []
-            
-            for pdf_file in pdf_files:
-                print(f"Processing: {pdf_file.name}")
-                
-                # Store PDF file path for later page rendering
-                self.pdf_files[pdf_file.name] = str(pdf_file)
-                
-                # Extract text with page information
-                pages = self.extract_text_from_pdf_with_pages(str(pdf_file))
-                self.pdf_pages[pdf_file.name] = pages
-                
-                # Combine all pages for vector database
-                full_text = "\n\n".join([f"Page {page_num}: {text}" for page_num, text in pages.items()])
-                
-                if full_text:
-                    all_texts.append(full_text)
-                    self.curriculum_docs.append({
-                        'filename': pdf_file.name,
-                        'content': full_text[:500] + "..." if len(full_text) > 500 else full_text,
-                        'pages': pages
-                    })
-            
-            if not all_texts:
-                print("No text could be extracted from PDF files!")
-                return False
-            
-            # Split text into smaller chunks with metadata
-            text_splitter = RecursiveCharacterTextSplitter(
-                chunk_size=500,  # Reduced from 1000
-                chunk_overlap=50,  # Reduced from 200
-                length_function=len,
-            )
-            
-            for i, text in enumerate(all_texts):
-                chunks = text_splitter.split_text(text)
-                for j, chunk in enumerate(chunks):
-                    # Add metadata to track which document and approximate page
-                    all_chunks_with_metadata.append({
-                        'text': chunk,
-                        'metadata': {
-                            'filename': pdf_files[i].name,
-                            'chunk_id': j,
-                            'source': 'curriculum'
-                        }
-                    })
-            
-            # Create embeddings
-            self.embeddings = HuggingFaceEmbeddings(
-                model_name="sentence-transformers/all-MiniLM-L6-v2"
-            )
-            
-            # Create vector database with metadata
-            texts = [chunk['text'] for chunk in all_chunks_with_metadata]
-            metadatas = [chunk['metadata'] for chunk in all_chunks_with_metadata]
-            
-            self.vector_db = Chroma.from_texts(
-                texts=texts,
-                embedding=self.embeddings,
-                metadatas=metadatas,
-                persist_directory="./chroma_db"
-            )
-            
-            print(f"Processed {len(pdf_files)} curriculum documents!")
-            return True
-            
-        except Exception as e:
-            print(f"Error processing curriculum: {str(e)}")
-            return False
-    
-    def create_qa_chain(self):
-        """Create the QA chain with custom prompts"""
-        if not self.vector_db or not self.llm:
-            return False
-            
-        # Better prompt template for more detailed responses
-        qa_template = """You are an expert programming instructor. Based on the curriculum context provided, answer the student's question in a clear and educational manner. Write a comprehensive paragraph that explains the concept thoroughly.
-
-Context: {context}
-
-Question: {question}
-
-Answer:"""
-
-        self.qa_chain = RetrievalQA.from_chain_type(
-            llm=self.llm,
-            chain_type="stuff",
-            retriever=self.vector_db.as_retriever(search_kwargs={"k": 3}),  # Increased for better context
-            chain_type_kwargs={
-                "prompt": PromptTemplate(
-                    template=qa_template,
-                    input_variables=["context", "question"]
-                )
-            }
-        )
-        
-        return True
-    
-    def find_relevant_pages(self, question: str, filename: str = None) -> List[Dict]:
-        """Find relevant pages for a given question"""
-        try:
-            # Search for relevant chunks
-            results = self.vector_db.similarity_search(question, k=5)  # Increased for better coverage
-            
-            relevant_pages = []
-            seen_pages = set()
-            
-            for result in results:
-                metadata = result.metadata
-                doc_filename = metadata.get('filename', '')
-                
-                # If filename is specified, only look in that file
-                if filename and doc_filename != filename:
-                    continue
-                
-                # Extract page information from chunk text
-                chunk_text = result.page_content
-                
-                # Look for page numbers in the chunk
-                page_matches = re.findall(r'Page (\d+):', chunk_text)
-                
-                for page_num in page_matches:
-                    page_key = f"{doc_filename}_page_{page_num}"
-                    if page_key not in seen_pages:
-                        seen_pages.add(page_key)
-                        
-                        # Get the actual page content
-                        if doc_filename in self.pdf_pages:
-                            page_content = self.pdf_pages[doc_filename].get(int(page_num), "")
-                            if page_content:
-                                relevant_pages.append({
-                                    'filename': doc_filename,
-                                    'page_number': int(page_num),
-                                    'content': page_content,
-                                    'relevance_score': len(chunk_text)  # Simple relevance metric
-                                })
-            
-            # Sort by relevance and return top results
-            relevant_pages.sort(key=lambda x: x['relevance_score'], reverse=True)
-            return relevant_pages[:3]  # Return top 3 most relevant pages
-            
-        except Exception as e:
-            print(f"Error finding relevant pages: {str(e)}")
-            return []
 
-def initialize_system():
-    """Initialize the curriculum assistant system"""
-    assistant = CurriculumAssistant()
-    
-    # Load LLM
-    if not assistant.load_llm():
-        return "❌ Failed to load language model", None, None
-    
-    # Process curriculum
-    if not assistant.process_curriculum("Slides"):
-        return "❌ Failed to process curriculum documents", None, None
-    
-    # Create QA chain
-    if not assistant.create_qa_chain():
-        return "❌ Failed to create QA chain", None, None
-    
-    return "✅ System initialized successfully!", assistant, assistant.curriculum_docs
+    def search(self, query):
+        # Find the most relevant chunk (page)
+        results = self.vector_db.similarity_search(query, k=1)
+        if not results:
+            return "No relevant page found.", None, None
+        result = results[0]
+        filename = result.metadata["filename"]
+        page_number = result.metadata["page_number"]
+        text = result.page_content
+        img = self.get_pdf_page_image(self.pdf_files[filename], page_number)
+        return text, img, f"{filename} - Page {page_number}"
 
-def ask_question(question: str, assistant: CurriculumAssistant):
-    """Ask a question and get answer with relevant pages"""
-    if not assistant or not assistant.qa_chain:
-        return "Please initialize the system first.", []
-    
-    try:
-        # Get answer from QA chain using invoke instead of run
-        answer = assistant.qa_chain.invoke({"query": question})
-        
-        # Find relevant pages
-        relevant_pages = assistant.find_relevant_pages(question)
-        
-        # Format page information and get page images
-        page_info = ""
-        page_images = []
-        
-        if relevant_pages:
-            page_info = "📄 **Relevant Pages Found:**\n\n"
-            for i, page in enumerate(relevant_pages, 1):
-                page_info += f"**{i}. {page['filename']} - Page {page['page_number']}**\n"
-                page_info += f"```\n{page['content'][:300]}...\n```\n\n"
-                
-                # Get page image
-                if page['filename'] in assistant.pdf_files:
-                    page_image = assistant.get_pdf_page_image(
-                        assistant.pdf_files[page['filename']], 
-                        page['page_number']
-                    )
-                    if page_image:
-                        page_images.append((page_image, f"{page['filename']} - Page {page['page_number']}"))
-                        print(f"Added page image for {page['filename']} page {page['page_number']}")
-                    else:
-                        print(f"Failed to get page image for {page['filename']} page {page['page_number']}")
-        else:
-            page_info = "No specific pages found for this question."
-        
-        # Format the complete response
-        full_response = f"## Answer\n\n{answer['result']}\n\n---\n\n{page_info}"
-        
-        return full_response, page_images
-        
-    except Exception as e:
-        error_msg = f"Error processing question: {str(e)}"
-        return error_msg, []
+# --- Gradio UI ---
+searcher = FastPDFSearch()
 
-# Initialize the system
-status, assistant, curriculum_docs = initialize_system()
+def gradio_search(query):
+    text, img, label = searcher.search(query)
+    if img:
+        return text, [(img, label)]
+    else:
+        return text, []
 
-# Create Gradio interface
-with gr.Blocks(title="Inclusive World Curriculum Assistant", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🎓 Inclusive World Curriculum Assistant")
-    gr.Markdown("An AI-powered assistant that answers questions about your curriculum and shows relevant slide pages.")
-    
+with gr.Blocks(title="Fast PDF Curriculum Search", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 📄 Fast PDF Curriculum Search\nAsk a question and see the most relevant slide page!")
     with gr.Row():
-        with gr.Column(scale=2):
-            # Status display
-            status_display = gr.Textbox(
-                value=status,
-                label="System Status",
-                interactive=False
-            )
-            
-            # Question input
-            question_input = gr.Textbox(
-                label="Ask a question about your curriculum",
-                placeholder="e.g., What are if statements? How do loops work?",
-                lines=3
-            )
-            
-            # Submit button
-            submit_btn = gr.Button("🔍 Get Answer", variant="primary")
-            
-            # Answer output
-            answer_output = gr.Markdown(
-                label="Answer with Relevant Pages",
-                value="Ask a question to get started!"
-            )
-        
-        with gr.Column(scale=1):
-            # Curriculum overview
-            gr.Markdown("### 📚 Curriculum Documents")
-            if curriculum_docs:
-                for doc in curriculum_docs:
-                    with gr.Accordion(f"📄 {doc['filename']}", open=False):
-                        gr.Markdown(f"**Preview:** {doc['content']}")
-            else:
-                gr.Markdown("No curriculum documents loaded.")
-    
-    # Page images display
-    with gr.Row():
-        gr.Markdown("### 📄 Relevant Slide Pages")
-        page_images_output = gr.Gallery(
-            label="PDF Pages",
-            show_label=True,
-            elem_id="gallery",
-            columns=2,
-            rows=2,
-            height="auto",
-            object_fit="contain"
-        )
-    
-    # Handle question submission
-    def process_question(question):
-        return ask_question(question, assistant)
-    
-    submit_btn.click(
-        fn=process_question,
-        inputs=[question_input],
-        outputs=[answer_output, page_images_output]
-    )
-    
-    # Handle Enter key in question input
-    question_input.submit(
-        fn=process_question,
-        inputs=[question_input],
-        outputs=[answer_output, page_images_output]
-    )
+        with gr.Column():
+            question = gr.Textbox(label="Ask a question", placeholder="e.g., What are for loops?", lines=2)
+            submit = gr.Button("🔍 Search")
+            answer = gr.Markdown(label="Relevant Page Text")
+        with gr.Column():
+            gallery = gr.Gallery(label="Relevant PDF Page", columns=1, rows=1, height="auto", object_fit="contain")
+    submit.click(fn=gradio_search, inputs=question, outputs=[answer, gallery])
+    question.submit(fn=gradio_search, inputs=question, outputs=[answer, gallery])
 
-# Launch the app
 if __name__ == "__main__":
     demo.launch()