Update app.py

app.py

@@ -1,283 +1,218 @@
-import gradio as gr
-from PyPDF2 import PdfReader
-from langchain.embeddings import HuggingFaceInferenceAPIEmbeddings
-from langchain.vectorstores import FAISS
-from huggingface_hub import InferenceClient
 import os
+import json
 import logging
-import traceback
-from datetime import datetime
-from typing import List, Dict, Tuple, Any
-import re
-
-# Configure logging
-logging.basicConfig(
-    level=logging.DEBUG,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-    handlers=[
-        logging.FileHandler(f'page_rag_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log'),
-        logging.StreamHandler()
-    ]
-)
-logger = logging.getLogger(__name__)
-
-class TextPreprocessor:
-    @staticmethod
-    def clean_text(text: str) -> str:
-        """Clean and normalize text content."""
-        # Remove multiple spaces
-        text = re.sub(r'\s+', ' ', text)
-        # Remove multiple newlines
-        text = re.sub(r'\n\s*\n', '\n\n', text)
-        # Normalize quotes
-        text = re.sub(r'["\'""]', '"', text)
-        # Remove header/footer artifacts
-        text = re.sub(r'^.*Page \d+.*$', '', text, flags=re.MULTILINE)
-        return text.strip()
-
-    @staticmethod
-    def extract_section_headers(text: str) -> List[str]:
-        """Extract potential section headers from text."""
-        # Simple header detection (can be enhanced based on document structure)
-        header_pattern = r'^(?:[A-Z][A-Za-z\s]{2,50}:?|(?:\d+\.){1,3}\s+[A-Z][A-Za-z\s]{2,50})$'
-        headers = re.findall(header_pattern, text, re.MULTILINE)
-        return headers
-
-def create_page_chunks(pdf_reader: PdfReader) -> List[Dict[str, Any]]:
-    """
-    Creates page-level chunks from PDF content.
-    """
-    page_chunks = []
-    preprocessor = TextPreprocessor()
+import shutil
+import gradio as gr
+from typing import List
+from tempfile import NamedTemporaryFile
+from huggingface_hub import InferenceClient
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain.docstore.document import Document
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+
+# Constants
+DOCUMENTS_FILE = "uploaded_documents.json"
+DEFAULT_MODEL = "@cf/meta/llama-2-7b-chat"
+HF_TOKEN = os.getenv("HF_API_TOKEN")  # Make sure to set this environment variable
+
+def get_embeddings():
+    return HuggingFaceEmbeddings(model_name="avsolatorio/GIST-Embedding-v0")
+
+def load_documents():
+    if os.path.exists(DOCUMENTS_FILE):
+        with open(DOCUMENTS_FILE, "r") as f:
+            return json.load(f)
+    return []
+
+def save_documents(documents):
+    with open(DOCUMENTS_FILE, "w") as f:
+        json.dump(documents, f)
+
+def load_document(file: NamedTemporaryFile) -> List[Document]:
+    """Loads and splits the document into pages using PyPDF."""
+    loader = PyPDFLoader(file.name)
+    return loader.load_and_split()
+
+def update_vectors(files):
+    if not files:
+        return "Please upload at least one file.", []
     
-    for page_num, page in enumerate(pdf_reader.pages, 1):
+    embed = get_embeddings()
+    uploaded_documents = load_documents()
+    total_chunks = 0
+    
+    all_data = []
+    for file in files:
         try:
-            page_text = page.extract_text()
-            if not page_text.strip():
+            data = load_document(file)
+            if not data:
                 continue
             
-            # Clean and preprocess text
-            cleaned_text = preprocessor.clean_text(page_text)
-            headers = preprocessor.extract_section_headers(cleaned_text)
+            all_data.extend(data)
+            total_chunks += len(data)
             
-            # Store full page as a chunk
-            page_chunks.append({
-                "content": cleaned_text,
-                "metadata": {
-                    "page_num": page_num,
-                    "section_headers": headers
-                }
-            })
+            if not any(doc["name"] == file.name for doc in uploaded_documents):
+                uploaded_documents.append({"name": file.name, "selected": True})
                 
         except Exception as e:
-            logger.error(f"Error processing page {page_num}: {str(e)}")
-            continue
-            
-    return page_chunks
-
-class RAGApplication:
-    def __init__(self, hf_api_key: str):
-        try:
-            self.hf_api_key = hf_api_key
-            self.vector_store = None
-            
-            logger.info("Initializing HuggingFace embeddings...")
-            self.embeddings = HuggingFaceInferenceAPIEmbeddings(
-                api_key=hf_api_key,
-                model_name="sentence-transformers/all-MiniLM-L6-v2"
-            )
-            
-            logger.info("Initializing HuggingFace client...")
-            self.client = InferenceClient(api_key=hf_api_key)
-            self.conversation_history = []
-            
-            # Initialize cache
-            self.query_cache = {}
-            
-            logger.info("RAGApplication initialized successfully")
-        except Exception as e:
-            logger.error(f"Error initializing RAGApplication: {str(e)}")
-            logger.error(f"Traceback: {traceback.format_exc()}")
-            raise
-
-        self.system_prompt = """You are a precise and accurate PDF summarization assistant. Your role is to:
-1. Provide accurate answers based solely on the provided context
-2. Maintain factual consistency and never hallucinate information
-3. Clearly indicate when information is not available in the context
-4. Use concise language and avoid unnecessary elaboration
-5. Maintain continuity with previous conversation when relevant
-
-Context: {context}
-
-Previous conversation:
-{conversation_history}
-
-Question: {question}
-
-Answer:"""
-
-    def process_pdf(self, file_path: str) -> str:
-        try:
-            logger.info(f"Starting PDF processing for file: {file_path}")
-            
-            if file_path is None or not os.path.exists(file_path):
-                return "Please upload a valid PDF file."
-
-            # Reset conversation history and cache
-            self.conversation_history = []
-            self.query_cache = {}
-            
-            pdf_reader = PdfReader(file_path)
-            
-            # Create page chunks
-            page_chunks = create_page_chunks(pdf_reader)
-            
-            # Create vector store
-            logger.info("Creating vector store...")
-            self.vector_store = FAISS.from_texts(
-                [chunk["content"] for chunk in page_chunks],
-                self.embeddings,
-                metadatas=[chunk["metadata"] for chunk in page_chunks]
-            )
-            
-            logger.info("Vector store created successfully")
-            return "PDF processed successfully!"
-            
-        except Exception as e:
-            logger.error(f"Error in PDF processing: {str(e)}")
-            return f"Error processing PDF: {str(e)}"
-
-    def retrieve_context(self, query: str, k: int = 3) -> str:
-        """
-        Retrieve relevant pages for the given query.
-        """
-        # Check query cache
-        cache_key = f"{query}_{k}"
-        if cache_key in self.query_cache:
-            return self.query_cache[cache_key]
-
-        # Get relevant pages
-        results = self.vector_store.similarity_search_with_score(query, k=k)
+            logging.error(f"Error processing file {file.name}: {str(e)}")
+    
+    if not all_data:
+        return "No valid data could be extracted from the uploaded files.", []
+    
+    try:
+        if os.path.exists("faiss_database"):
+            database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+            database.add_documents(all_data)
+        else:
+            database = FAISS.from_documents(all_data, embed)
+        database.save_local("faiss_database")
         
-        # Combine context while preserving document structure
-        context = []
-        for doc, score in results:
-            context_str = f"[Page {doc.metadata['page_num']}"
-            
-            if doc.metadata.get('section_headers'):
-                context_str += f", Section: {doc.metadata['section_headers'][0]}"
-                
-            context_str += f"]: {doc.page_content}"
-            context.append(context_str)
-                
-        final_context = "\n\n".join(context)
+        save_documents(uploaded_documents)
+        return f"Vector store updated successfully. Processed {total_chunks} chunks.", uploaded_documents
         
-        # Cache the result
-        self.query_cache[cache_key] = final_context
-        return final_context
-
-    def generate_response(self, message: str, history: List[Tuple[str, str]]) -> str:
-        try:
-            logger.info(f"Generating response for message: {message}")
-            
-            if not self.vector_store:
-                return "Please upload and process a PDF first."
-
-            query = message.strip()
-            if not query:
-                return "Please enter a question."
-
-            # Get relevant context
-            context = self.retrieve_context(query)
-            
-            # Format conversation history
-            conversation_history = "\n".join([
-                f"Q: {q}\nA: {a}" for q, a in history[-3:] if q and a
-            ])
-
-            # Create prompt
-            prompt = self.system_prompt.format(
-                context=context,
-                conversation_history=conversation_history,
-                question=query
-            )
-
-            # Generate response using Mistral
-            logger.info("Generating response using Mistral...")
-            response = ""
-            try:
-                for message in self.client.chat_completion(
-                    model="mistralai/Mistral-Nemo-Instruct-2407",
-                    messages=[
-                        {"role": "system", "content": prompt},
-                        {"role": "user", "content": query}
-                    ],
-                    max_tokens=10000,
-                    stream=True,
-                ):
-                    response += message.choices[0].delta.content
-                logger.info("Response generated successfully")
-            except Exception as e:
-                logger.error(f"Error in chat completion: {str(e)}")
-                raise
+    except Exception as e:
+        return f"Error updating vector store: {str(e)}", []
 
-            return response
-        except Exception as e:
-            error_msg = f"Error generating response: {str(e)}"
-            logger.error(error_msg)
-            logger.error(f"Traceback: {traceback.format_exc()}")
-            return error_msg
+def delete_documents(selected_docs):
+    if not selected_docs:
+        return "No documents selected for deletion.", []
+    
+    uploaded_documents = load_documents()
+    embed = get_embeddings()
+    
+    if os.path.exists("faiss_database"):
+        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+        
+        docs_to_keep = []
+        for doc in database.docstore._dict.values():
+            if doc.metadata.get("source") not in selected_docs:
+                docs_to_keep.append(doc)
+        
+        if not docs_to_keep:
+            shutil.rmtree("faiss_database")
+        else:
+            new_database = FAISS.from_documents(docs_to_keep, embed)
+            new_database.save_local("faiss_database")
+        
+        uploaded_documents = [doc for doc in uploaded_documents if doc["name"] not in selected_docs]
+        save_documents(uploaded_documents)
+        
+        return f"Deleted documents: {', '.join(selected_docs)}", uploaded_documents
+    
+    return "No documents to delete.", []
 
-def create_gradio_interface():
+def get_response(query, temperature=0.2):
+    if not query.strip():
+        return "Please enter a question."
+    
+    uploaded_documents = load_documents()
+    selected_docs = [doc["name"] for doc in uploaded_documents if doc["selected"]]
+    
+    if not selected_docs:
+        return "Please select at least one document to search through."
+    
+    embed = get_embeddings()
+    if not os.path.exists("faiss_database"):
+        return "No documents available. Please upload PDF documents first."
+    
+    database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+    
+    # Filter documents
+    filtered_docs = []
+    for doc in database.docstore._dict.values():
+        if isinstance(doc, Document) and doc.metadata.get("source") in selected_docs:
+            filtered_docs.append(doc)
+    
+    if not filtered_docs:
+        return "No relevant information found in the selected documents."
+    
+    filtered_db = FAISS.from_documents(filtered_docs, embed)
+    retriever = filtered_db.as_retriever(search_kwargs={"k": 5})
+    relevant_docs = retriever.get_relevant_documents(query)
+    
+    context_str = "\n".join([doc.page_content for doc in relevant_docs])
+    
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant that provides accurate answers based on the given context."},
+        {"role": "user", "content": f"Context:\n{context_str}\n\nQuestion: {query}\n\nProvide a comprehensive answer based only on the given context."}
+    ]
+    
+    client = InferenceClient(DEFAULT_MODEL, token=HF_TOKEN)
+    
     try:
-        logger.info("Creating Gradio interface...")
-        api_key = os.getenv("HF_API_KEY")
-        rag = RAGApplication(hf_api_key=api_key)
+        response = client.chat_completion(
+            messages=messages,
+            max_tokens=1000,
+            temperature=temperature,
+            top_p=0.9,
+        )
+        return response.choices[0].message.content
+    except Exception as e:
+        return f"Error generating response: {str(e)}"
+
+def create_interface():
+    with gr.Blocks(title="PDF Question Answering System") as app:
+        gr.Markdown("# PDF Question Answering System")
         
-        with gr.Blocks() as demo:
-            gr.Markdown("# PDF Question Answering System")
-            
-            with gr.Row():
-                pdf_input = gr.File(
-                    label="Upload PDF",
+        with gr.Row():
+            with gr.Column():
+                files = gr.File(
+                    label="Upload PDF Documents",
                     file_types=[".pdf"],
-                    type="filepath"
+                    multiple=True
                 )
-                process_button = gr.Button("Process PDF")
-                status_output = gr.Textbox(label="Status", interactive=False)
+                upload_button = gr.Button("Upload and Process")
+                
+            with gr.Column():
+                doc_status = gr.Textbox(label="Status", interactive=False)
+                doc_list = gr.Checkboxgroup(
+                    label="Available Documents",
+                    choices=[],
+                    interactive=True
+                )
+                delete_button = gr.Button("Delete Selected Documents")
+        
+        with gr.Row():
+            with gr.Column():
+                question = gr.Textbox(label="Ask a question about the documents")
+                temperature = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.2,
+                    step=0.1,
+                    label="Temperature"
+                )
+                submit_button = gr.Button("Submit Question")
             
-            process_button.click(
-                fn=rag.process_pdf,
-                inputs=[pdf_input],
-                outputs=[status_output]
-            )
-
-            chat_interface = gr.ChatInterface(
-                fn=rag.generate_response,
-                title="Chat with your PDF",
-                description="Upload a PDF and ask questions about its contents.",
-                theme="soft",
-                examples=[
-                    "What is the main topic of this document?",
-                    "Can you summarize the key points?",
-                    "What are the main conclusions?",
-                ],
-            )
+            with gr.Column():
+                answer = gr.Textbox(label="Answer", interactive=False)
         
-        logger.info("Gradio interface created successfully")
-        return demo
-    except Exception as e:
-        logger.error(f"Error creating Gradio interface: {str(e)}")
-        logger.error(f"Traceback: {traceback.format_exc()}")
-        raise
+        # Event handlers
+        upload_button.click(
+            fn=update_vectors,
+            inputs=[files],
+            outputs=[doc_status, doc_list]
+        )
+        
+        delete_button.click(
+            fn=delete_documents,
+            inputs=[doc_list],
+            outputs=[doc_status, doc_list]
+        )
+        
+        submit_button.click(
+            fn=get_response,
+            inputs=[question, temperature],
+            outputs=[answer]
+        )
+    
+    return app
 
 if __name__ == "__main__":
-    try:
-        logger.info("Starting application...")
-        demo = create_gradio_interface()
-        logger.info("Launching Gradio interface...")
-        demo.launch()
-    except Exception as e:
-        logger.error(f"Application failed to start: {str(e)}")
-        logger.error(f"Traceback: {traceback.format_exc()}")
-        raise
+    app = create_interface()
+    app.launch()