Update app.py

app.py

@@ -1,10 +1,12 @@
 import pandas as pd
 import torch
 from sentence_transformers import SentenceTransformer, util
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import gradio as gr
 import json
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import faiss
 import spaces
+
 # Ensure you have GPU support
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
@@ -12,41 +14,48 @@ device = 'cuda' if torch.cuda.is_available() else 'cpu'
 df = pd.read_csv('RBDx10kstats.csv')
 df['embedding'] = df['embedding'].apply(json.loads)  # Convert JSON string back to list
 
-# Convert embeddings to tensor for efficient retrieval
-embeddings = torch.tensor(df['embedding'].tolist(), device=device)
+# Convert embeddings to a numpy array
+embeddings = np.array(df['embedding'].tolist(), dtype='float32')
+
+# Setup FAISS
+index = faiss.IndexFlatL2(embeddings.shape[1])  # dimension should match the embedding size
+index.add(embeddings)
 
 # Load the Sentence Transformer model
-model = SentenceTransformer('all-MiniLM-L6-v2', device=device)
+sentence_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2', device=device)
+
+# Load the LLaMA model for response generation
+llama_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
+llama_model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf").to(device)
 
-# Load the ai model for response generation
-ai_tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2-large")
-ai_model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2-large").to(device)
+# Load the summarization model
+summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=0 if device == 'cuda' else -1)
 
-# Define the function to find the most relevant document
-@spaces.GPU(duration=120)
+# Define the function to find the most relevant document using FAISS
 def retrieve_relevant_doc(query):
-    query_embedding = model.encode(query, convert_to_tensor=True, device=device)
-    similarities = util.pytorch_cos_sim(query_embedding, embeddings)[0]
-    best_match_idx = torch.argmax(similarities).item()
+    query_embedding = sentence_model.encode(query, convert_to_tensor=False)
+    _, indices = index.search(np.array([query_embedding]), k=1)
+    best_match_idx = indices[0][0]
     return df.iloc[best_match_idx]['Abstract']
 
 # Define the function to generate a response
-@spaces.GPU(duration=120)
 def generate_response(query):
     relevant_doc = retrieve_relevant_doc(query)
+    if len(relevant_doc) > 512:  # Truncate long documents
+        relevant_doc = summarizer(relevant_doc, max_length=150, min_length=50, do_sample=False)[0]['summary_text']
     input_text = f"Document: {relevant_doc}\n\nQuestion: {query}\n\nAnswer:"
-    inputs = ai_tokenizer(input_text, return_tensors="pt").to(device)
-    outputs = ai_model.generate(inputs["input_ids"], max_length=1024)
-    response = ai_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    inputs = llama_tokenizer(input_text, return_tensors="pt").to(device)
+    outputs = llama_model.generate(inputs["input_ids"], max_length=150)
+    response = llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response
 
 # Create a Gradio interface
 iface = gr.Interface(
     fn=generate_response,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your query here..."),
+    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter your query here..."),
     outputs="text",
     title="RAG Chatbot",
-    description="This chatbot retrieves relevant documents based on your query and generates responses using ai models."
+    description="This chatbot retrieves relevant documents based on your query and generates responses using LLaMA."
 )
 
 # Launch the Gradio interface