Update app.py

app.py

@@ -1,43 +1,15 @@
+# app.py
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
 import torch
-import time  # Added for timing
-# New imports
+import time
+import os
+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
 
-# Document processing function
-def process_documents(files):
-    """Process PDF/TXT files into vector embeddings"""
-    documents = []
-    for file_path in 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())
-    
-    # Split documents into chunks
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=512,
-        chunk_overlap=50
-    )
-    texts = text_splitter.split_documents(documents)
-    
-    # Create embeddings
-    embeddings = HuggingFaceEmbeddings(
-        model_name="BAAI/bge-small-en-v1.5"
-    )
-    
-    # Create vector store
-    vectorstore = FAISS.from_documents(texts, embeddings)
-    return vectorstore
-
-
+# Configuration
 MODEL_CONFIG = {
     "phi-3": {
         "model_name": "microsoft/phi-3-mini-4k-instruct",
@@ -46,8 +18,7 @@ MODEL_CONFIG = {
     "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|>
-"""
+{message}<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
     }
 }
 
@@ -62,33 +33,78 @@ class ChatModel:
     def __init__(self):
         self.models = {}
         self.tokenizers = {}
-        self.vectorstore = None  # Add vectorstore reference
+        self.vectorstore = None
 
-    # Add this new method
-    def update_vectorstore(self, files):
-        """Process uploaded files and update vectorstore"""
-        if files:
-            self.vectorstore = process_documents(files)
+    def load_model(self, model_name):
+        if model_name not in self.models:
+            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
+
+    def process_documents(self, files, progress=gr.Progress()):
+        """Process uploaded documents into vector embeddings"""
+        try:
+            progress(0, desc="Starting document processing")
+            documents = []
+            
+            # 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())
+            
+            # Split documents
+            progress(0.3, desc="Processing documents")
+            text_splitter = RecursiveCharacterTextSplitter(
+                chunk_size=512,
+                chunk_overlap=50
+            )
+            texts = text_splitter.split_documents(documents)
+            
+            # Create embeddings
+            progress(0.6, desc="Generating embeddings")
+            embeddings = HuggingFaceEmbeddings(
+                model_name="BAAI/bge-small-en-v1.5"
+            )
+            
+            # Create vector store
+            progress(0.8, desc="Building vector database")
+            self.vectorstore = FAISS.from_documents(texts, embeddings)
+            
+            return "✅ Documents processed successfully! Ready for queries."
+        
+        except Exception as e:
+            return f"❌ Error processing documents: {str(e)}"
 
-    # Modify existing generate method
     def generate(self, message, model_name, history):
         start_time = time.time()
+        self.load_model(model_name)
+        config = MODEL_CONFIG[model_name]
         
         # Retrieve relevant context
         context = ""
         if self.vectorstore:
             docs = self.vectorstore.similarity_search(message, k=3)
-            context = "\n".join([d.page_content for d in docs])
+            context = "\n\n".join([d.page_content for d in docs])
         
-        self.load_model(model_name)
-        config = MODEL_CONFIG[model_name]
-        
-        # Update prompt with context
+        # Format prompt with context
         prompt = config["template"].format(
-            message=f"Context: {context}\n\nQuestion: {message}"
+            message=f"Context:\n{context}\n\nQuestion: {message}"
         )
         
-        # Create pipeline
+        # Generate response
         pipe = pipeline(
             "text-generation",
             model=self.models[model_name],
@@ -110,6 +126,7 @@ class ChatModel:
         
         return response, elapsed_time, tokens_per_sec
 
+# Initialize model handler
 model_handler = ChatModel()
 
 def chat(message, history, model_choice):
@@ -123,33 +140,39 @@ def chat(message, history, model_choice):
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🚀 LLM Chatbot with RAG & Performance Metrics")
     
-    # Add document upload section
-    with gr.Row():
-        file_upload = gr.File(
-            label="Upload Documents (PDF/TXT)",
-            file_count="multiple",
-            file_types=[".pdf", ".txt"],
-            type="filepath"
-        )
     with gr.Row():
         model_choice = gr.Dropdown(
             choices=["phi-3", "llama3-8b"],
             label="Select Model",
             value="phi-3"
         )
-    chatbot = gr.Chatbot(height=400)
-    msg = gr.Textbox(label="Message", placeholder="Type here...")
+    
+    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...")
+    
     with gr.Row():
         submit_btn = gr.Button("Send", variant="primary")
-        clear_btn = gr.ClearButton([msg, chatbot])
-
-    msg.submit(chat, [msg, chatbot, model_choice], chatbot)
-    submit_btn.click(chat, [msg, chatbot, model_choice], chatbot)
+        clear_btn = gr.ClearButton([msg, chatbot, file_upload])
 
+    # Event handlers
     file_upload.upload(
-        fn=model_handler.update_vectorstore,
+        fn=model_handler.process_documents,
         inputs=file_upload,
-        outputs=None
+        outputs=status,
+        show_progress="full"
     )
+    
+    msg.submit(chat, [msg, chatbot, model_choice], chatbot)
+    submit_btn.click(chat, [msg, chatbot, model_choice], chatbot)
 
 demo.launch()