File size: 2,625 Bytes
43fda8b 6f8702e 43fda8b 6f8702e 43fda8b 960f72f 43fda8b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import os
import streamlit as st
import pandas as pd
from datasets import load_from_disk
from transformers import AutoTokenizer, TFAutoModel
from constant import DRGUS_STR_LIST
if DRGUS_STR_LIST:
Drugs = DRGUS_STR_LIST.split(',')
Drugs = [drug.strip() for drug in Drugs]
model_ckpt = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
model = TFAutoModel.from_pretrained(model_ckpt, from_pt=True)
def cls_pooling(model_output):
return model_output.last_hidden_state[:, 0]
def get_embeddings(text_list):
encoded_input = tokenizer(
text_list, padding=True, truncation=True, return_tensors="tf"
)
encoded_input = {k: v for k, v in encoded_input.items()}
model_output = model(**encoded_input)
return cls_pooling(model_output)
embeddings_dataset = load_from_disk("data")
embeddings_dataset.add_faiss_index(column="embeddings")
def recommendations(question):
question_embedding = get_embeddings([question]).numpy()
scores, samples = embeddings_dataset.get_nearest_examples(
"embeddings", question_embedding, k=5
)
samples_df = pd.DataFrame.from_dict(samples)
samples_df["scores"] = scores
samples_df.sort_values("scores", ascending=False, inplace=True,ignore_index=True)
return samples_df[['drugName', 'review', 'scores']]
# Create Streamlit app
st.title("Call on Doc Drug Recommendation System")
st.markdown(
"""
<style>
#MainMenu {visibility: hidden;}
footer {visibility: hidden;}
</style>
""",
unsafe_allow_html=True
)
# Allow users to select a default question or input their own
st.sidebar.title("Choose or Enter a Question:")
selection_type = st.sidebar.radio("Select type:", ("Select Default", "Enter Custom"))
if selection_type == "Select Default":
selected_question = st.sidebar.selectbox("Select a question", Drugs)
if st.sidebar.button("Show Recommendations"):
recommendation_result = recommendations(selected_question)
st.header(f"Top 5 Recommended Drugs for '{selected_question}':")
st.table(recommendation_result)
else:
default_question = "I've acne problem"
custom_question = st.sidebar.text_input("Enter your question:", default_question)
if st.sidebar.button("Get Recommendations"):
if custom_question:
custom_recommendation_result = recommendations(custom_question)
st.header("Top 5 Recommended Drugs for Your Question:")
st.table(custom_recommendation_result)
else:
st.warning("Please enter a question to get recommendations.") |