Update app.py

app.py

@@ -12,6 +12,7 @@ from langchain.prompts import PromptTemplate
 from langchain_community.llms import HuggingFacePipeline
 from langchain_community.document_loaders import PyPDFLoader, TextLoader, Docx2txtLoader
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import login
 
 # Configure logging
 logging.basicConfig(
@@ -20,129 +21,128 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
+# Constants
 MODEL_NAME = "meta-llama/Llama-2-7b-chat-hf"
 UPLOAD_FOLDER = "uploaded_docs"
+EMBEDDING_MODEL = "intfloat/multilingual-e5-large"
 
-class DocumentManager:
-    """Class to manage document uploads and processing."""
+class RAGSystem:
+    """Main RAG system class."""
     
     def __init__(self):
         self.upload_folder = UPLOAD_FOLDER
         if os.path.exists(self.upload_folder):
             shutil.rmtree(self.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']
+        
+        # Initialize components
+        self.embeddings = None
+        self.vector_store = None
+        self.qa_chain = None
         self.documents = []
         
-    def validate_file(self, file_path, file_size):
+        # Initialize embeddings once
+        self.initialize_embeddings()
+        
+    def initialize_embeddings(self):
+        """Initialize embedding model."""
+        try:
+            self.embeddings = HuggingFaceEmbeddings(
+                model_name=EMBEDDING_MODEL,
+                model_kwargs={'device': 'cuda' if torch.cuda.is_available() else 'cpu'}
+            )
+        except Exception as e:
+            logger.error(f"Error initializing embeddings: {str(e)}")
+            raise
+
+    def validate_file(self, file_path: str, file_size: int) -> bool:
+        """Validate uploaded file."""
         if file_size > self.max_file_size:
             raise ValueError(f"File size exceeds {self.max_file_size // 1024 // 1024}MB limit")
         
         ext = os.path.splitext(file_path)[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()
+            raise ValueError(f"Unsupported format. Supported: {', '.join(self.supported_formats)}")
+        return True
+
+    def process_file(self, file: gr.File) -> List:
+        """Process a single file and return documents."""
         try:
+            file_path = file.name
+            file_size = os.path.getsize(file_path)
+            self.validate_file(file_path, file_size)
+            
+            # Copy file to upload directory
+            filename = os.path.basename(file_path)
+            save_path = os.path.join(self.upload_folder, filename)
+            shutil.copy2(file_path, save_path)
+            
+            # Load documents based on file type
+            ext = os.path.splitext(file_path)[1].lower()
             if ext == '.pdf':
-                loader = PyPDFLoader(file_path)
+                loader = PyPDFLoader(save_path)
             elif ext == '.txt':
-                loader = TextLoader(file_path)
-            elif ext == '.docx':
-                loader = Docx2txtLoader(file_path)
-            else:
-                raise ValueError(f"Unsupported file format: {ext}")
+                loader = TextLoader(save_path)
+            else:  # .docx
+                loader = Docx2txtLoader(save_path)
                 
             documents = loader.load()
             for doc in documents:
                 doc.metadata.update({
-                    'source': os.path.basename(file_path),
+                    'source': filename,
                     'type': 'uploaded'
                 })
             return documents
             
         except Exception as e:
-            logger.error(f"Error loading {file_path}: {str(e)}")
+            logger.error(f"Error processing {file_path}: {str(e)}")
             raise
 
-    def process_upload(self, files: List[gr.File]) -> str:
-        if not files:
-            return "No files uploaded"
-            
-        current_files = len(os.listdir(self.upload_folder))
-        if current_files + len(files) > self.max_files:
-            return f"Maximum number of documents ({self.max_files}) exceeded"
-            
-        processed_files = []
-        for file in files:
-            try:
-                file_path = file.name
-                file_size = os.path.getsize(file_path)
-                
-                self.validate_file(file_path, file_size)
-                
-                # Copy file to upload folder
-                filename = os.path.basename(file_path)
-                save_path = os.path.join(self.upload_folder, filename)
-                shutil.copy2(file_path, save_path)
-                
-                docs = self.load_document(save_path)
-                self.documents.extend(docs)
-                processed_files.append(filename)
-                
-            except Exception as e:
-                logger.error(f"Error processing {file_path}: {str(e)}")
-                return f"Error processing {os.path.basename(file_path)}: {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."""
+    def update_vector_store(self, new_documents: List):
+        """Update vector store with new documents."""
         try:
-            if not documents:
-                raise ValueError("No documents provided for initialization")
-                
-            # Initialize text splitter
+            # Process documents
             text_splitter = RecursiveCharacterTextSplitter(
                 chunk_size=500,
                 chunk_overlap=50,
                 separators=["\n\n", "\n", ". ", " ", ""]
             )
+            chunks = text_splitter.split_documents(new_documents)
             
-            # 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 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)
+                
+        except Exception as e:
+            logger.error(f"Error updating vector store: {str(e)}")
+            raise
+
+    def initialize_llm(self):
+        """Initialize the language model and QA chain."""
+        try:
+            # Get Hugging Face token
+            hf_token = os.environ.get('HUGGINGFACE_TOKEN')
+            if not hf_token:
+                raise ValueError("Please set HUGGINGFACE_TOKEN environment variable")
             
-            # Create vector store
-            self.vector_store = FAISS.from_documents(chunks, self.embeddings)
+            # Login to Hugging Face
+            login(token=hf_token)
             
-            # Initialize LLM pipeline
-            tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            # Initialize model and tokenizer
+            tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
             model = AutoModelForCausalLM.from_pretrained(
-                self.model_name,
+                MODEL_NAME,
                 torch_dtype=torch.float16,
                 device_map="auto"
             )
             
+            # Create pipeline
             pipe = pipeline(
                 "text-generation",
                 model=model,
@@ -154,7 +154,7 @@ class RAGSystem:
             
             llm = HuggingFacePipeline(pipeline=pipe)
             
-            # Create prompt template
+            # Create QA chain
             prompt_template = """
             Context: {context}
             
@@ -169,7 +169,6 @@ class RAGSystem:
                 input_variables=["context", "question"]
             )
             
-            # Set up QA chain
             self.qa_chain = RetrievalQA.from_chain_type(
                 llm=llm,
                 chain_type="stuff",
@@ -178,16 +177,44 @@ class RAGSystem:
                 chain_type_kwargs={"prompt": PROMPT}
             )
             
-            self.is_initialized = True
-            return "System initialized successfully"
+        except Exception as e:
+            logger.error(f"Error initializing LLM: {str(e)}")
+            raise
+
+    def process_upload(self, files: List[gr.File]) -> str:
+        """Process uploaded files and initialize/update the system."""
+        if not files:
+            return "Please select files to upload."
+            
+        try:
+            current_files = len(os.listdir(self.upload_folder))
+            if current_files + len(files) > self.max_files:
+                return f"Maximum number of documents ({self.max_files}) exceeded"
+            
+            # Process each file
+            processed_files = []
+            new_documents = []
+            for file in files:
+                documents = self.process_file(file)
+                new_documents.extend(documents)
+                processed_files.append(os.path.basename(file.name))
+            
+            # Update vector store with new documents
+            self.update_vector_store(new_documents)
+            self.documents.extend(new_documents)
+            
+            # Initialize LLM if not already initialized
+            if self.qa_chain is None:
+                self.initialize_llm()
+            
+            return f"Successfully processed and initialized: {', '.join(processed_files)}"
             
         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:
+        if not self.qa_chain:
             return {"error": "System not initialized. Please upload documents first."}
             
         try:
@@ -211,25 +238,13 @@ class RAGSystem:
             logger.error(f"Error generating response: {str(e)}")
             return {"error": str(e)}
 
-# Initialize RAG system
+# Initialize 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 or "Maximum" 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."""
+def process_query(message: str, history: List) -> List:
+    """Process user query and return updated history."""
     try:
-        if not rag_system.is_initialized:
+        if not rag_system.qa_chain:
             return history + [(message, "Please upload documents first.")]
             
         response = rag_system.generate_response(message)
@@ -257,17 +272,14 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     with gr.Row():
         # Sidebar for document upload
         with gr.Column(scale=1):
-            # Removed gr.Box() and using styling directly
-            with gr.Group(
-                elem_classes="container",
-            ):
+            with gr.Group():
                 gr.HTML("""
                     <div style="padding: 1rem; border: 1px solid #e5e7eb; border-radius: 0.5rem; background-color: white;">
-                    <h3 style="margin-top: 0;">📁 Document Upload</h3>
+                    <h3 style="margin-top: 0;">📁 Upload Documents</h3>
                 """)
                 file_output = gr.File(
                     file_count="multiple",
-                    label="Upload Documents",
+                    label="Select Files",
                     elem_id="file-upload"
                 )
                 gr.HTML("""
@@ -277,12 +289,11 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                         <p>• Supported: PDF, TXT, DOCX</p>
                     </div>
                 """)
-                upload_button = gr.Button("📤 Upload and Initialize", variant="primary")
                 system_output = gr.Textbox(
-                    label="System Status",
+                    label="Status",
                     interactive=False
                 )
-                gr.HTML("</div>") # Closing the styled container
+                gr.HTML("</div>")
         
         # Main chat area
         with gr.Column(scale=3):
@@ -326,7 +337,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     """)
 
     # Add custom CSS
-    css = """
+    demo.css = """
         .container {
             border-radius: 0.5rem;
             margin: 0.5rem;
@@ -335,11 +346,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             margin-bottom: 1rem;
         }
     """
-    demo.css = css
     
     # Set up event handlers
-    upload_button.click(
-        process_file_upload,
+    file_output.upload(
+        rag_system.process_upload,
         inputs=[file_output],
         outputs=[system_output]
     )