Spaces:
Runtime error
Runtime error
| import streamlit as st | |
| import transformers | |
| from annotated_text import annotated_text | |
| ENTITY_TO_COLOR = { | |
| 'PER': '#8ef', | |
| 'LOC': '#faa', | |
| 'ORG': '#afa', | |
| 'MISC': '#fea', | |
| } | |
| def get_pipe(model_name): | |
| model = transformers.AutoModelForTokenClassification.from_pretrained(model_name) | |
| tokenizer = transformers.AutoTokenizer.from_pretrained(model_name) | |
| pipe = transformers.pipeline("token-classification", | |
| model=model, | |
| tokenizer=tokenizer, | |
| aggregation_strategy="simple") | |
| return pipe | |
| def parse_text(text, prediction): | |
| start = 0 | |
| parsed_text = [] | |
| for p in prediction: | |
| parsed_text.append(text[start:p["start"]]) | |
| parsed_text.append((p["word"], p["entity_group"], ENTITY_TO_COLOR[p["entity_group"]])) | |
| start = p["end"] | |
| parsed_text.append(text[start:]) | |
| return parsed_text | |
| st.set_page_config(page_title="Named Entity Recognition") | |
| st.title("Named Entity Recognition") | |
| st.write("Type text into the text box and then press 'Predict' to get the named entities.") | |
| option = st.selectbox('Select model', | |
| ("dslim/bert-base-NER", | |
| "dslim/bert-large-NER", | |
| "Davlan/bert-base-multilingual-cased-ner-hrl")) | |
| default_text = """Mark Zuckerberg has unveiled a new VR headset, the Quest Pro, at an online event held for developers. With a price tag of $1,499, Quest Pro is almost four times the price of Meta's current headset, the Quest 2, which starts at $399. It boasts thinner lenses, a curved battery around the head strap at the back, and controllers which self-track. The headset also allows users to see their real environment around the periphery of the screen. Quest Pro has mixed reality capabilities, unlike its predecessor - meaning digital content can be viewed overlaid on the real world. | |
| Meta boss Mark Zuckerberg said mixed reality was "the next major step for VR". | |
| Gartner analyst Tuong Nguyen said the high launch price made it a device more suited to "high-end, enthusiast and potentially enterprise users" than mass-market. Microsoft CEO Satya Nadella also announced that the office platform Windows 365 will be available on it, as Meta seeks to position mixed and virtual reality as a work tool as well as a form of entertainment.""" | |
| text = st.text_area('Enter text here:', value=default_text) | |
| st.write('Model used for prediction:', option) | |
| submit = st.button('Predict') | |
| with st.spinner("Loading model..."): | |
| pipe = get_pipe(model_name=option) | |
| if (submit and len(text.strip()) > 0) or len(text.strip()) > 0: | |
| prediction = pipe(text) | |
| parsed_text = parse_text(text, prediction) | |
| st.header("Prediction:") | |
| annotated_text(*parsed_text) | |