import pandas as pd
import torch
from sentence_transformers import SentenceTransformer, util
import gradio as gr
import json
from transformers import AutoTokenizer, AutoModelForCausalLM
import spaces
# 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['Abstract'] = df['Abstract'].apply(json.loads)  # Convert JSON string back to list

# Convert embeddings to tensor for efficient retrieval
embeddings = torch.tensor(df['Abstract'].tolist(), device=device)

# Load the Sentence Transformer model
model = SentenceTransformer('all-MiniLM-L6-v2', device=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)

# 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, 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
@spaces.GPU(duration=120)
def generate_response(query):
    relevant_doc = retrieve_relevant_doc(query)
    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)
    return response

# Create a Gradio interface
iface = gr.Interface(
    fn=generate_response,
    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 and generates responses using ai models."
)

# Launch the Gradio interface
iface.launch()