Update app.py

app.py

@@ -2,6 +2,29 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
 import torch
 import time  # Added for timing
+# New imports
+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):
+    documents = []
+    for file in files:
+        if file.name.endswith(".pdf"):
+            loader = PyPDFLoader(file.name)
+        elif file.name.endswith(".txt"):
+            loader = TextLoader(file.name)
+        documents.extend(loader.load())
+    
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=50)
+    texts = text_splitter.split_documents(documents)
+    
+    embeddings = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en-v1.5")
+    vectorstore = FAISS.from_documents(texts, embeddings)
+    return vectorstore
+
 
 MODEL_CONFIG = {
     "phi-3": {
@@ -45,7 +68,13 @@ class ChatModel:
             self.models[model_name] = model
             self.tokenizers[model_name] = tokenizer
 
-    def generate(self, message, model_name, history):
+    def generate(self, message, model_name, history, vectorstore=None):
+        # RAG context retrieval
+        if vectorstore:
+            docs = vectorstore.similarity_search(message, k=3)
+            context = "\n".join([d.page_content for d in docs])
+            message = f"Context: {context}\n\nQuestion: {message}"
+
         start_time = time.time()  # Start timing
         self.load_model(model_name)
         config = MODEL_CONFIG[model_name]
@@ -86,7 +115,12 @@ def chat(message, history, model_choice):
         return [(message, f"Error: {str(e)}")]
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚀 LLM Chatbot with Performance Metrics")
+    gr.Markdown("# 🚀 LLM Chatbot with RAG & Performance Metrics")
+    
+    # Add document upload section
+    with gr.Row():
+        file_output = gr.File(label="Upload Documents", file_count="multiple", 
+                            file_types=[".pdf", ".txt"], max_size=10)
     with gr.Row():
         model_choice = gr.Dropdown(
             choices=["phi-3", "llama3-8b"],