Create app.py

app.py

@@ -0,0 +1,309 @@
+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, TextLoader, Docx2txtLoader
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"
+UPLOAD_FOLDER = "uploaded_docs"
+
+class DocumentManager:
+    """Class to manage document uploads and processing."""
+    
+    def __init__(self):
+        self.upload_folder = UPLOAD_FOLDER
+        os.makedirs(self.upload_folder, exist_ok=True)
+        self.max_files = 5
+        self.max_file_size = 10 * 1024 * 1024  # 10 MB
+        self.supported_formats = ['.pdf', '.txt', '.docx']
+        self.documents = []
+        
+    def validate_file(self, file):
+        if os.path.getsize(file.name) > self.max_file_size:
+            raise ValueError(f"File size exceeds {self.max_file_size // 1024 // 1024}MB limit")
+        
+        ext = os.path.splitext(file.name)[1].lower()
+        if ext not in self.supported_formats:
+            raise ValueError(f"Unsupported file format. Supported formats: {', '.join(self.supported_formats)}")
+            
+    def load_document(self, file_path: str) -> List:
+        ext = os.path.splitext(file_path)[1].lower()
+        try:
+            if ext == '.pdf':
+                loader = PyPDFLoader(file_path)
+            elif ext == '.txt':
+                loader = TextLoader(file_path)
+            elif ext == '.docx':
+                loader = Docx2txtLoader(file_path)
+            else:
+                raise ValueError(f"Unsupported file format: {ext}")
+                
+            documents = loader.load()
+            for doc in documents:
+                doc.metadata.update({
+                    'source': os.path.basename(file_path),
+                    'type': 'uploaded'
+                })
+            return documents
+            
+        except Exception as e:
+            logger.error(f"Error loading {file_path}: {str(e)}")
+            raise
+
+    def process_upload(self, files: List) -> str:
+        if len(os.listdir(self.upload_folder)) + len(files) > self.max_files:
+            raise ValueError(f"Maximum number of documents ({self.max_files}) exceeded")
+            
+        processed_files = []
+        for file in files:
+            try:
+                self.validate_file(file)
+                save_path = os.path.join(self.upload_folder, file.name)
+                file.save(save_path)
+                docs = self.load_document(save_path)
+                self.documents.extend(docs)
+                processed_files.append(file.name)
+            except Exception as e:
+                logger.error(f"Error processing {file.name}: {str(e)}")
+                return f"Error processing {file.name}: {str(e)}"
+                
+        return f"Successfully processed files: {', '.join(processed_files)}"
+
+class RAGSystem:
+    """Main RAG system class."""
+    
+    def __init__(self, model_name: str = MODEL_NAME):
+        self.model_name = model_name
+        self.document_manager = DocumentManager()
+        self.embeddings = None
+        self.vector_store = None
+        self.qa_chain = None
+        self.is_initialized = False
+        
+    def initialize_system(self, documents: List = None):
+        """Initialize RAG system with provided documents."""
+        try:
+            if not documents:
+                raise ValueError("No documents provided for initialization")
+                
+            # Initialize text splitter
+            text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=500,
+                chunk_overlap=50,
+                separators=["\n\n", "\n", ". ", " ", ""]
+            )
+            
+            # Process documents
+            chunks = text_splitter.split_documents(documents)
+            
+            # Initialize embeddings
+            self.embeddings = HuggingFaceEmbeddings(
+                model_name="intfloat/multilingual-e5-large",
+                model_kwargs={'device': 'cuda' if torch.cuda.is_available() else 'cpu'}
+            )
+            
+            # Create vector store
+            self.vector_store = FAISS.from_documents(chunks, self.embeddings)
+            
+            # Initialize LLM pipeline
+            tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            model = AutoModelForCausalLM.from_pretrained(
+                self.model_name,
+                torch_dtype=torch.float16,
+                device_map="auto"
+            )
+            
+            pipe = pipeline(
+                "text-generation",
+                model=model,
+                tokenizer=tokenizer,
+                max_new_tokens=512,
+                temperature=0.1,
+                device_map="auto"
+            )
+            
+            llm = HuggingFacePipeline(pipeline=pipe)
+            
+            # Create prompt template
+            prompt_template = """
+            Context: {context}
+            
+            Based on the context above, please provide a clear and concise answer to the following question.
+            If the information is not in the context, explicitly state so.
+            
+            Question: {question}
+            """
+            
+            PROMPT = PromptTemplate(
+                template=prompt_template,
+                input_variables=["context", "question"]
+            )
+            
+            # Set up QA chain
+            self.qa_chain = RetrievalQA.from_chain_type(
+                llm=llm,
+                chain_type="stuff",
+                retriever=self.vector_store.as_retriever(search_kwargs={"k": 4}),
+                return_source_documents=True,
+                chain_type_kwargs={"prompt": PROMPT}
+            )
+            
+            self.is_initialized = True
+            return "System initialized successfully"
+            
+        except Exception as e:
+            logger.error(f"Error during system initialization: {str(e)}")
+            return f"Error: {str(e)}"
+
+    def generate_response(self, question: str) -> Dict:
+        """Generate response for a given question."""
+        if not self.is_initialized:
+            return {"error": "System not initialized. Please upload documents first."}
+            
+        try:
+            result = self.qa_chain({"query": question})
+            
+            response = {
+                'answer': result['result'],
+                'sources': []
+            }
+            
+            for doc in result['source_documents']:
+                source = {
+                    'title': doc.metadata.get('source', 'Unknown'),
+                    'content': doc.page_content[:200] + "..." if len(doc.page_content) > 200 else doc.page_content
+                }
+                response['sources'].append(source)
+            
+            return response
+            
+        except Exception as e:
+            logger.error(f"Error generating response: {str(e)}")
+            return {"error": str(e)}
+
+# Initialize RAG system
+rag_system = RAGSystem()
+
+def process_file_upload(files):
+    """Handle file uploads and system initialization."""
+    try:
+        upload_result = rag_system.document_manager.process_upload(files)
+        if "Error" in upload_result:
+            return upload_result
+            
+        init_result = rag_system.initialize_system(rag_system.document_manager.documents)
+        return f"{upload_result}\n{init_result}"
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+def process_query(message, history):
+    """Process user query and generate response."""
+    try:
+        if not rag_system.is_initialized:
+            return history + [(message, "Please upload documents first.")]
+            
+        response = rag_system.generate_response(message)
+        if "error" in response:
+            return history + [(message, f"Error: {response['error']}")]
+            
+        answer = response['answer']
+        sources = set([source['title'] for source in response['sources']])
+        if sources:
+            answer += "\n\n📚 Sources:\n" + "\n".join([f"• {source}" for source in sources])
+            
+        return history + [(message, answer)]
+    except Exception as e:
+        return history + [(message, f"Error: {str(e)}")]
+
+# Create Gradio interface
+demo = gr.Blocks(css="div.gradio-container {background-color: #f0f2f6}")
+
+with demo:
+    gr.HTML("""
+        <div style="text-align: center; max-width: 800px; margin: 0 auto; padding: 20px;">
+            <h1 style="color: #2d333a;">🤖 Easy RAG</h1>
+            <p style="color: #4a5568;">
+                A simple and powerful RAG system for your documents
+            </p>
+        </div>
+    """)
+    
+    with gr.Row():
+        file_output = gr.File(
+            file_count="multiple",
+            label="Upload Documents (PDF, TXT, DOCX - Max 5 files, 10MB each)"
+        )
+    
+    upload_button = gr.Button("Upload and Initialize")
+    system_output = gr.Textbox(label="System Status")
+    
+    chatbot = gr.Chatbot(
+        show_label=False,
+        container=True,
+        height=400,
+        show_copy_button=True
+    )
+    
+    with gr.Row():
+        message = gr.Textbox(
+            placeholder="Ask a question about your documents...",
+            show_label=False,
+            container=False,
+            scale=8
+        )
+        clear = gr.Button("🗑️ Clear", size="sm", scale=1)
+    
+    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 Easy RAG</h3>
+                <p style="color: #666; font-size: 14px;">
+                    A powerful RAG system that lets you query your documents using:
+                </p>
+                <ul style="list-style: none; color: #666; font-size: 14px;">
+                    <li>🔹 LLM: 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;">
+                    Based on original work by <a href="https://www.linkedin.com/in/camilo-vega-169084b1/" 
+                    target="_blank" style="color: #2196F3; text-decoration: none;">Camilo Vega</a>
+                </p>
+            </div>
+        </div>
+    """)
+    
+    # Set up event handlers
+    upload_button.click(
+        process_file_upload,
+        inputs=[file_output],
+        outputs=[system_output]
+    )
+    
+    message.submit(
+        process_query,
+        inputs=[message, chatbot],
+        outputs=[chatbot]
+    )
+    
+    clear.click(lambda: None, None, chatbot)
+
+demo.launch()