import streamlit as st
from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
from datasets import load_dataset
import torch

# Load the dataset
dataset = load_dataset("pubmed_qa", split="test")

# Initialize RAG components
tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="default", use_dummy_dataset=True)
model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq")

# Function to get the answer to a medical query
def get_medical_answer(query):
    # Encode the query to retrieve relevant documents
    inputs = tokenizer(query, return_tensors="pt")
    input_ids = inputs["input_ids"]
    
    # Retrieve relevant documents
    docs = retriever(input_ids=input_ids, return_tensors="pt")
    
    # Generate the answer from the model
    generated_ids = model.generate(input_ids=input_ids, context_input_ids=docs["context_input_ids"], 
                                   context_attention_mask=docs["context_attention_mask"])
    
    # Decode the generated answer
    generated_answer = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return generated_answer

# Streamlit UI
st.title("Medical QA Assistant")
st.write("Ask any medical question, and I will answer it based on PubMed papers!")

# Input text box for queries
query = st.text_input("Enter your medical question:")

if query:
    with st.spinner("Searching for the answer..."):
        answer = get_medical_answer(query)
    st.write(f"Answer: {answer}")