AI-health-assistant / app (1).py
Haniyamsohail's picture
Upload app (1).py
07406f3 verified
raw
history blame
954 Bytes
import streamlit as st
from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
# Initialize the model
tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")
retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", use_dummy_dataset=True)
# Streamlit UI
st.title("AI Health Assistant (RAG-based)")
def get_answer_rag(question):
inputs = tokenizer(question, return_tensors="pt")
retrieved_docs = retriever.retrieve(inputs['input_ids'], top_k=3)
outputs = model.generate(input_ids=inputs['input_ids'], context_input_ids=retrieved_docs['input_ids'])
answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
return answer
# Ask the user for a health-related question
question = st.text_input("Ask a health-related question:")
if question:
answer = get_answer_rag(question)
st.write(f"Answer: {answer}")