import streamlit as st
from transformers import pipeline
from PIL import Image
import pandas as pd
import matplotlib.pyplot as plt

labels = ['AS-yes', 'AS-no', 'AS-no_information']

# Change `transformersbook` to your Hub username
model_id = "sarahflan/distilbert-base-uncased-finetuned-AS_sentences"
classifier = pipeline("text-classification", model=model_id)

st.title("Aortic Stenosis sentence classifier ")

sentence = st.text_input('Echocardiogram demonstrated a heavily calcified aortic valve with limited leaflet mobility. Findings are consistent with severe aortic stenosis.','Patient has AS')

preds = classifier(sentence, return_all_scores=True)

preds_df = pd.DataFrame(preds[0], columns=['score'])
preds_df['labels'] = labels

sorted_preds_df = preds_df.sort_values(by='score', ascending=False).head(20)

fig, ax = plt.subplots()
ax.bar(sorted_preds_df['labels'], 100 * sorted_preds_df['score'], color='C0')
ax.set_title(f'"{sentence}"')
ax.set_ylabel("Class probability (%)")
ax.tick_params(axis='x', labelrotation=90)

st.pyplot(fig)