Update app.py

app.py

@@ -3,22 +3,42 @@ import gradio as gr
 import torch
 import time
 import os
+import logging
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
 from langchain_community.document_loaders import PyPDFLoader, TextLoader
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 # Configuration
 MODEL_CONFIG = {
     "phi-3": {
         "model_name": "microsoft/phi-3-mini-4k-instruct",
-        "template": "<|user|>\n{message}<|end|>\n<|assistant|>"
+        "template": """<|user|>
+Using the following context, please answer the question. If the context doesn't contain relevant information, say so.
+
+Context:
+{context}
+
+Question: {question}<|end|>
+<|assistant|>
+Let me help answer your question based on the provided context."""
     },
     "llama3-8b": {
         "model_name": "NousResearch/Meta-Llama-3-8B-Instruct",
         "template": """<|begin_of_text|><|start_header_id|>user<|end_header_id|>
-{message}<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
+Using the following context, please answer the question. If the context doesn't contain relevant information, say so.
+
+Context:
+{context}
+
+Question: {question}<|eot_id|>
+<|start_header_id|>assistant<|end_header_id|>
+Let me help answer your question based on the provided context."""
     }
 }
 
@@ -34,131 +54,142 @@ class ChatModel:
         self.models = {}
         self.tokenizers = {}
         self.vectorstore = {}
+        self.embeddings = HuggingFaceEmbeddings(
+            model_name="BAAI/bge-small-en-v1.5"
+        )
 
     def load_model(self, model_name):
+        """Load and cache the model and tokenizer"""
         if model_name not in self.models:
+            logger.info(f"Loading model: {model_name}")
             config = MODEL_CONFIG[model_name]
-            tokenizer = AutoTokenizer.from_pretrained(config["model_name"])
-            tokenizer.pad_token = tokenizer.eos_token
-            model = AutoModelForCausalLM.from_pretrained(
-                config["model_name"],
-                quantization_config=bnb_config,
-                device_map="auto",
-                torch_dtype=torch.float16,
-            )
-            self.models[model_name] = model
-            self.tokenizers[model_name] = tokenizer
-
-
+            try:
+                tokenizer = AutoTokenizer.from_pretrained(config["model_name"])
+                tokenizer.pad_token = tokenizer.eos_token
+                model = AutoModelForCausalLM.from_pretrained(
+                    config["model_name"],
+                    quantization_config=bnb_config,
+                    device_map="auto",
+                    torch_dtype=torch.float16,
+                )
+                self.models[model_name] = model
+                self.tokenizers[model_name] = tokenizer
+                logger.info(f"Successfully loaded model: {model_name}")
+            except Exception as e:
+                logger.error(f"Error loading model {model_name}: {str(e)}")
+                raise
 
     def load_vector_store(self, store_name):
-        """Cache vector stores in memory"""
-        
+        """Load and cache vector stores"""
         if store_name not in self.vectorstore:
-            embeddings = HuggingFaceEmbeddings(
-                model_name="BAAI/bge-small-en-v1.5"
-            )
-            self.vectorstore[store_name] = FAISS.load_local(
-                f"vector_stores_index/{store_name}", 
-                embeddings,
-                allow_dangerous_deserialization=True
-            )
+            logger.info(f"Loading vector store: {store_name}")
+            try:
+                self.vectorstore[store_name] = FAISS.load_local(
+                    f"vector_stores_index/{store_name}", 
+                    self.embeddings,
+                    allow_dangerous_deserialization=True
+                )
+                # Verify vector store content
+                self.check_vectorstore(store_name)
+                logger.info(f"Successfully loaded vector store: {store_name}")
+            except Exception as e:
+                logger.error(f"Error loading vector store {store_name}: {str(e)}")
+                raise
         return self.vectorstore[store_name]
 
+    def check_vectorstore(self, store_name):
+        """Verify vector store content"""
+        try:
+            vectorstore = self.vectorstore[store_name]
+            sample_query = "test query"
+            docs = vectorstore.similarity_search(sample_query, k=1)
+            logger.info(f"Sample document from {store_name}: {docs[0].page_content[:200]}...")
+        except Exception as e:
+            logger.error(f"Error checking vector store {store_name}: {str(e)}")
+            raise
 
-
-    def process_documents(self, files, progress=gr.Progress()):
-        """Process uploaded documents into vector embeddings"""
+    def generate(self, message, model_name, vector_store_name, history):
+        """Generate response using RAG"""
+        start_time = time.time()
         try:
-            progress(0, desc="Starting document processing")
-            documents = []
+            # Load model and vector store
+            self.load_model(model_name)
+            vectorstore = self.load_vector_store(vector_store_name)
+            config = MODEL_CONFIG[model_name]
             
-            # Load documents
-            for file_path in progress.tqdm(files, desc="Loading files"):
-                if file_path.endswith(".pdf"):
-                    loader = PyPDFLoader(file_path)
-                elif file_path.endswith(".txt"):
-                    loader = TextLoader(file_path)
-                else:
-                    continue
-                documents.extend(loader.load())
+            # Retrieve relevant context
+            logger.info(f"Retrieving context for query: {message}")
+            docs = vectorstore.similarity_search(message, k=3)
+            context = "\n\n".join([d.page_content for d in docs])
+            logger.info(f"Retrieved context: {context[:200]}...")
             
-            # Split documents
-            progress(0.3, desc="Processing documents")
-            text_splitter = RecursiveCharacterTextSplitter(
-                chunk_size=512,
-                chunk_overlap=50
+            # Format prompt
+            prompt = config["template"].format(
+                context=context,
+                question=message
             )
-            texts = text_splitter.split_documents(documents)
             
-            # Create embeddings
-            progress(0.6, desc="Generating embeddings")
-            embeddings = HuggingFaceEmbeddings(
-                model_name="BAAI/bge-small-en-v1.5"
+            # Generate response
+            pipe = pipeline(
+                "text-generation",
+                model=self.models[model_name],
+                tokenizer=self.tokenizers[model_name],
+                max_new_tokens=384,
+                temperature=0.3,  # Lower temperature for more focused responses
+                top_p=0.9,
+                repetition_penalty=1.1,
+                do_sample=True,
+                return_full_text=False
             )
             
-            # Create vector store
-            progress(0.8, desc="Building vector database")
-            self.vectorstore = FAISS.from_documents(texts, embeddings)
+            response = pipe(prompt)[0]['generated_text']
             
-            return "✅ Documents processed successfully! Ready for queries."
-        
-        except Exception as e:
-            return f"❌ Error processing documents: {str(e)}"
+            # Calculate metrics
+            elapsed_time = time.time() - start_time
+            tokens = len(self.tokenizers[model_name].encode(response))
+            tokens_per_sec = tokens / elapsed_time if elapsed_time > 0 else 0
+            
+            logger.info(f"Generated response in {elapsed_time:.2f}s")
+            return response, elapsed_time, tokens_per_sec
 
-    def generate(self, message, model_name, vector_store_name, history):
-        start_time = time.time()
-        self.load_model(model_name)
-        vectorstore = self.load_vector_store(vector_store_name)
-        config = MODEL_CONFIG[model_name]
-        
-        # Retrieve relevant context
-        context = ""
-        # if vectorstore:
-        docs = vectorstore.similarity_search(message, k=3)
-        context = "\n\n".join([d.page_content for d in docs])
-        
-        # Format prompt with context
-        prompt = config["template"].format(
-            message=f"Context:\n{context}\n\nQuestion: {message}"
-        )
-        
-        # Generate response
-        pipe = pipeline(
-            "text-generation",
-            model=self.models[model_name],
-            tokenizer=self.tokenizers[model_name],
-            max_new_tokens=384,
-            temperature=0.7,
-            top_p=0.9,
-            repetition_penalty=1.1,
-            do_sample=True,
-            return_full_text=False
-        )
-        
-        response = pipe(prompt)[0]['generated_text']
-        
-        # Calculate metrics
-        elapsed_time = time.time() - start_time
-        tokens = len(self.tokenizers[model_name].encode(response))
-        tokens_per_sec = tokens / elapsed_time if elapsed_time > 0 else 0
-        
-        return response, elapsed_time, tokens_per_sec
+        except Exception as e:
+            logger.error(f"Error in generate: {str(e)}")
+            raise
 
 # Initialize model handler
 model_handler = ChatModel()
 
 def chat(message, history, model_choice, vector_store_choice):
-    print("vector_store_choice: ", vector_store_choice)
+    """Chat interface function"""
+    logger.info(f"Received message: {message}")
+    logger.info(f"Using model: {model_choice}")
+    logger.info(f"Using vector store: {vector_store_choice}")
+    
     try:
-        response, response_time, token_speed = model_handler.generate(message, model_choice, vector_store_choice, history)
-        formatted_response = f"{response}\n\n⏱️ Response Time: {response_time:.2f}s | 🚀 Speed: {token_speed:.2f} tokens/s"
+        response, response_time, token_speed = model_handler.generate(
+            message, 
+            model_choice, 
+            vector_store_choice, 
+            history
+        )
+        
+        # Format response with metrics
+        formatted_response = (
+            f"{response}\n\n"
+            f"⏱️ Response Time: {response_time:.2f}s | "
+            f"🚀 Speed: {token_speed:.2f} tokens/s"
+        )
+        
         return [(message, formatted_response)]
+    
     except Exception as e:
-        return [(message, f"Error: {str(e)}")]
+        logger.error(f"Error in chat: {str(e)}")
+        error_message = f"Error: {str(e)}\n\nPlease try again or contact support if the issue persists."
+        return [(message, error_message)]
 
+# Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚀 LLM Chatbot with RAG & Performance Metrics")
+    gr.Markdown("# 🚀 Enhanced RAG Chatbot with Performance Metrics")
     
     with gr.Row():
         model_choice = gr.Dropdown(
@@ -173,31 +204,21 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         )
 
     with gr.Row():
-        with gr.Column(scale=1):
-            file_upload = gr.File(
-                label="Upload Documents (PDF/TXT)",
-                file_count="multiple",
-                file_types=[".pdf", ".txt"],
-                type="filepath"
-            )
-            status = gr.Textbox(label="Processing Status", interactive=False)
         with gr.Column(scale=3):
             chatbot = gr.Chatbot(height=500)
-            msg = gr.Textbox(label="Message", placeholder="Type your question here...")
+            msg = gr.Textbox(
+                label="Message",
+                placeholder="Type your question here...",
+                scale=4
+            )
     
     with gr.Row():
         submit_btn = gr.Button("Send", variant="primary")
-        clear_btn = gr.ClearButton([msg, chatbot, file_upload])
+        clear_btn = gr.ClearButton([msg, chatbot])
 
     # Event handlers
-    file_upload.upload(
-        fn=model_handler.process_documents,
-        inputs=file_upload,
-        outputs=status,
-        show_progress="full"
-    )
-    
     msg.submit(chat, [msg, chatbot, model_choice, vector_store_choice], chatbot)
     submit_btn.click(chat, [msg, chatbot, model_choice, vector_store_choice], chatbot)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()