from transformers import AutoModel, AutoTokenizer import streamlit as st st.header("MeshTagger 🔖") threshold = st.sidebar.slider("Threshold", value=0.5, min_value=0.0, max_value=1.0) display_probabilities = st.sidebar.checkbox("Display probabilities") if "model" not in st.session_state: with st.spinner("Loading model and tokenizer..."): st.session_state["tokenizer"] = AutoTokenizer.from_pretrained( "Wellcome/WellcomeBertMesh" ) st.session_state["model"] = AutoModel.from_pretrained( "Wellcome/WellcomeBertMesh", trust_remote_code=True ) model = st.session_state["model"] tokenizer = st.session_state["tokenizer"] text = st.text_area("", value="This text is about Malaria", height=400) inputs = tokenizer([text], padding="max_length") outputs = model(**inputs)[0] if display_probabilities: data = [ (model.id2label[label_id], label_prob.item()) for label_id, label_prob in enumerate(outputs) if label_prob > threshold ] st.table(data) else: for label_id, label_prob in enumerate(outputs): if label_prob > threshold: st.button(model.id2label[label_id]) st.markdown( "Examples from The Wellcome Trust are taken from data that are publishes openly at [360 Giving](https://data.threesixtygiving.org/). They are published under a [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/) license." )