Create app.py

app.py

@@ -0,0 +1,349 @@
+import os
+import logging
+from typing import List, Dict
+import torch
+import gradio as gr
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.chains import RetrievalQA
+from langchain.prompts import PromptTemplate
+from langchain.llms import HuggingFacePipeline
+from langchain_community.document_loaders import (
+    PyPDFLoader,
+    Docx2txtLoader,
+    CSVLoader,
+    UnstructuredFileLoader
+)
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+import spaces
+import tempfile
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Constants
+MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"
+SUPPORTED_FORMATS = [".pdf", ".docx", ".doc", ".csv", ".txt"]
+
+class DocumentLoader:
+    """Enhanced document loader supporting multiple file formats."""
+    
+    @staticmethod
+    def load_file(file_path: str) -> List:
+        """Load a single file based on its extension."""
+        ext = os.path.splitext(file_path)[1].lower()
+        try:
+            if ext == '.pdf':
+                loader = PyPDFLoader(file_path)
+            elif ext in ['.docx', '.doc']:
+                loader = Docx2txtLoader(file_path)
+            elif ext == '.csv':
+                loader = CSVLoader(file_path)
+            else:  # fallback for txt and other text files
+                loader = UnstructuredFileLoader(file_path)
+            
+            documents = loader.load()
+            
+            # Add metadata
+            for doc in documents:
+                doc.metadata.update({
+                    'title': os.path.basename(file_path),
+                    'type': 'document',
+                    'format': ext[1:],
+                    'language': 'auto'
+                })
+            
+            logger.info(f"Successfully loaded {file_path}")
+            return documents
+            
+        except Exception as e:
+            logger.error(f"Error loading {file_path}: {str(e)}")
+            raise
+
+class RAGSystem:
+    """Enhanced RAG system with dynamic document loading."""
+    
+    def __init__(self, model_name: str = MODEL_NAME):
+        self.model_name = model_name
+        self.embeddings = None
+        self.vector_store = None
+        self.qa_chain = None
+        self.tokenizer = None
+        self.model = None
+        self.is_initialized = False
+    
+    def initialize_model(self):
+        """Initialize the base model and tokenizer."""
+        try:
+            logger.info("Initializing language model...")
+            
+            # Initialize embeddings
+            self.embeddings = HuggingFaceEmbeddings(
+                model_name="intfloat/multilingual-e5-large",
+                model_kwargs={'device': 'cuda' if torch.cuda.is_available() else 'cpu'},
+                encode_kwargs={'normalize_embeddings': True}
+            )
+            
+            # Initialize model and tokenizer
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model_name,
+                trust_remote_code=True
+            )
+            
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_name,
+                torch_dtype=torch.float16,
+                trust_remote_code=True,
+                device_map="auto"
+            )
+            
+            # Create generation pipeline
+            pipe = pipeline(
+                "text-generation",
+                model=self.model,
+                tokenizer=self.tokenizer,
+                max_new_tokens=512,
+                temperature=0.1,
+                top_p=0.95,
+                repetition_penalty=1.15,
+                device_map="auto"
+            )
+            
+            self.llm = HuggingFacePipeline(pipeline=pipe)
+            self.is_initialized = True
+            
+            logger.info("Model initialization completed")
+            
+        except Exception as e:
+            logger.error(f"Error during model initialization: {str(e)}")
+            raise
+
+    def process_documents(self, files: List[tempfile._TemporaryFileWrapper]) -> None:
+        """Process uploaded documents and update the vector store."""
+        try:
+            documents = []
+            for file in files:
+                docs = DocumentLoader.load_file(file.name)
+                documents.extend(docs)
+            
+            if not documents:
+                raise ValueError("No documents were successfully loaded.")
+            
+            # Process documents
+            text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=800,
+                chunk_overlap=200,
+                separators=["\n\n", "\n", ". ", " ", ""],
+                length_function=len
+            )
+            
+            chunks = text_splitter.split_documents(documents)
+            
+            # Create or update vector store
+            if self.vector_store is None:
+                self.vector_store = FAISS.from_documents(chunks, self.embeddings)
+            else:
+                self.vector_store.add_documents(chunks)
+            
+            # Initialize QA chain
+            prompt_template = """
+            Context: {context}
+            
+            Based on the provided context, please answer the following question clearly and concisely.
+            If the information is not in the context, please say so explicitly.
+            
+            Question: {question}
+            """
+            
+            PROMPT = PromptTemplate(
+                template=prompt_template,
+                input_variables=["context", "question"]
+            )
+            
+            self.qa_chain = RetrievalQA.from_chain_type(
+                llm=self.llm,
+                chain_type="stuff",
+                retriever=self.vector_store.as_retriever(
+                    search_kwargs={"k": 6}
+                ),
+                return_source_documents=True,
+                chain_type_kwargs={"prompt": PROMPT}
+            )
+            
+            logger.info(f"Successfully processed {len(documents)} documents")
+            
+        except Exception as e:
+            logger.error(f"Error processing documents: {str(e)}")
+            raise
+
+    def generate_response(self, question: str) -> Dict:
+        """Generate response for a given question."""
+        if not self.is_initialized or self.qa_chain is None:
+            return {
+                'answer': "Please upload some documents first before asking questions.",
+                'sources': []
+            }
+        
+        try:
+            result = self.qa_chain({"query": question})
+            
+            response = {
+                'answer': result['result'],
+                'sources': []
+            }
+            
+            for doc in result['source_documents']:
+                source = {
+                    'title': doc.metadata.get('title', 'Unknown'),
+                    'content': doc.page_content[:200] + "..." if len(doc.page_content) > 200 else doc.page_content,
+                    'metadata': doc.metadata
+                }
+                response['sources'].append(source)
+            
+            return response
+            
+        except Exception as e:
+            logger.error(f"Error generating response: {str(e)}")
+            raise
+
+@spaces.GPU(duration=60)
+def process_response(user_input: str, chat_history: List, files: List) -> tuple:
+    """Process user input and generate response."""
+    try:
+        if not rag_system.is_initialized:
+            rag_system.initialize_model()
+        
+        if files and (rag_system.vector_store is None):
+            rag_system.process_documents(files)
+        
+        response = rag_system.generate_response(user_input)
+        
+        # Clean and format response
+        answer = response['answer']
+        if "Answer:" in answer:
+            answer = answer.split("Answer:")[-1].strip()
+        
+        # Format sources
+        sources = set([source['title'] for source in response['sources'][:3]])
+        if sources:
+            answer += "\n\n📚 Sources consulted:\n" + "\n".join([f"• {source}" for source in sources])
+        
+        chat_history.append((user_input, answer))
+        return chat_history
+        
+    except Exception as e:
+        logger.error(f"Error in process_response: {str(e)}")
+        error_message = f"Sorry, an error occurred: {str(e)}"
+        chat_history.append((user_input, error_message))
+        return chat_history
+
+# Initialize RAG system
+logger.info("Initializing RAG system...")
+try:
+    rag_system = RAGSystem()
+    logger.info("RAG system created successfully")
+except Exception as e:
+    logger.error(f"Failed to create RAG system: {str(e)}")
+    raise
+
+# Create Gradio interface
+try:
+    logger.info("Creating Gradio interface...")
+    with gr.Blocks(css="div.gradio-container {background-color: #f0f2f6}") as demo:
+        gr.HTML("""
+            <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
+                <h1 style="color: #2d333a;">📚 DocumentGPT</h1>
+                <p style="color: #4a5568;">
+                    Your AI Assistant for Document Analysis and Q&A
+                </p>
+            </div>
+        """)
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                files = gr.Files(
+                    label="Upload Your Documents",
+                    file_types=SUPPORTED_FORMATS,
+                    file_count="multiple"
+                )
+                gr.HTML("""
+                    <div style="font-size: 0.9em; color: #666; margin-top: 0.5em;">
+                        Supported formats: PDF, DOCX, CSV, TXT
+                    </div>
+                """)
+
+        chatbot = gr.Chatbot(
+            show_label=False,
+            container=True,
+            height=500,
+            bubble_full_width=True,
+            show_copy_button=True,
+            scale=2
+        )
+        
+        with gr.Row():
+            message = gr.Textbox(
+                placeholder="💭 Ask me anything about your documents...",
+                show_label=False,
+                container=False,
+                scale=8,
+                autofocus=True
+            )
+            clear = gr.Button("🗑️ Clear", size="sm", scale=1)
+
+        # Instructions
+        gr.HTML("""
+            <div style="background-color: #f8f9fa; padding: 15px; border-radius: 10px; margin: 20px 0;">
+                <h3 style="color: #2d333a; margin-bottom: 10px;">🔍 How to use:</h3>
+                <ol style="color: #666; margin-left: 20px;">
+                    <li>Upload one or more documents (PDF, DOCX, CSV, or TXT)</li>
+                    <li>Wait for the documents to be processed</li>
+                    <li>Ask questions about your documents</li>
+                    <li>View sources used in the responses</li>
+                </ol>
+            </div>
+        """)
+
+        # Footer with credits
+        gr.HTML("""
+            <div style="text-align: center; max-width: 800px; margin: 20px auto; padding: 20px;
+                        background-color: #f8f9fa; border-radius: 10px;">
+                <div style="margin-bottom: 15px;">
+                    <h3 style="color: #2d333a;">⚡ About this assistant</h3>
+                    <p style="color: #666; font-size: 14px;">
+                        This application uses RAG (Retrieval Augmented Generation) technology combining:
+                    </p>
+                    <ul style="list-style: none; color: #666; font-size: 14px;">
+                        <li>🔹 LLM Engine: Llama-2-7b-chat-hf</li>
+                        <li>🔹 Embeddings: multilingual-e5-large</li>
+                        <li>🔹 Vector Store: FAISS</li>
+                    </ul>
+                </div>
+                <div style="border-top: 1px solid #ddd; padding-top: 15px;">
+                    <p style="color: #666; font-size: 14px;">
+                        Created by <a href="https://www.linkedin.com/in/camilo-vega-169084b1/" 
+                        target="_blank" style="color: #2196F3; text-decoration: none;">Camilo Vega</a>,
+                        AI Professor and Solutions Consultant 🤖
+                    </p>
+                </div>
+            </div>
+        """)
+
+        # Configure event handlers
+        def submit(user_input, chat_history, files):
+            return process_response(user_input, chat_history, files)
+            
+        message.submit(submit, [message, chatbot, files], [chatbot])
+        clear.click(lambda: None, None, chatbot)
+
+    logger.info("Gradio interface created successfully")
+    demo.launch()
+
+except Exception as e:
+    logger.error(f"Error in Gradio interface creation: {str(e)}")
+    raise