Haniyamsohail commited on
Commit
a00553e
·
verified ·
1 Parent(s): 74a8eae

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +23 -0
app.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
3
+
4
+ # Initialize the model
5
+ tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
6
+ model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq")
7
+ retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", use_dummy_dataset=True)
8
+
9
+ # Streamlit UI
10
+ st.title("AI Health Assistant (RAG-based)")
11
+
12
+ def get_answer_rag(question):
13
+ inputs = tokenizer(question, return_tensors="pt")
14
+ retrieved_docs = retriever.retrieve(inputs['input_ids'], top_k=3)
15
+ outputs = model.generate(input_ids=inputs['input_ids'], context_input_ids=retrieved_docs['input_ids'])
16
+ answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
17
+ return answer
18
+
19
+ # Ask the user for a health-related question
20
+ question = st.text_input("Ask a health-related question:")
21
+ if question:
22
+ answer = get_answer_rag(question)
23
+ st.write(f"Answer: {answer}")