Update app.py

app.py

@@ -3,33 +3,40 @@ import torch
 from sentence_transformers import SentenceTransformer, util
 import gradio as gr
 import json
-import spaces
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# Ensure you have GPU support
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 # Load the CSV file with embeddings
 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())
+embeddings = torch.tensor(df['embedding'].tolist(), device=device)
+
+# Load the Sentence Transformer model
+model = SentenceTransformer('all-MiniLM-L6-v2', device=device)
 
-# Load the same Sentence Transformer model
-model = SentenceTransformer('all-MiniLM-L6-v2')
+# Load the LLaMA model for response generation
+llama_tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
+llama_model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct").to(device)
 
 # Define the function to find the most relevant document
-@spaces.GPU(duration=120)
 def retrieve_relevant_doc(query):
-    query_embedding = model.encode(query, convert_to_tensor=True)
+    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()
     return df.iloc[best_match_idx]['Abstract']
 
-# Define the function to generate a response (for simplicity, echo the retrieved doc)
-@spaces.GPU(duration=120)
+# Define the function to generate a response
 def generate_response(query):
     relevant_doc = retrieve_relevant_doc(query)
-    # Here you could use a more sophisticated language model to generate a response
-    # For now, we will just return the relevant document as the response
-    return relevant_doc
+    input_text = f"Document: {relevant_doc}\n\nQuestion: {query}\n\nAnswer:"
+    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(
@@ -37,7 +44,7 @@ iface = gr.Interface(
     inputs=gr.Textbox(lines=2, placeholder="Enter your query here..."),
     outputs="text",
     title="RAG Chatbot",
-    description="This chatbot retrieves relevant documents based on your query."
+    description="This chatbot retrieves relevant documents based on your query and generates responses using LLaMA."
 )
 
 # Launch the Gradio interface