Spaces:
Runtime error
Runtime error
| import streamlit as st | |
| from transformers import AutoModel, AutoTokenizer | |
| 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]) | |