feat: added models choice

app.py

@@ -1,6 +1,6 @@
 import streamlit as st
-from annotated_text import annotated_text
 import transformers
+from annotated_text import annotated_text
 
 ENTITY_TO_COLOR = {
     'PER': '#8ef',
@@ -9,14 +9,18 @@ ENTITY_TO_COLOR = {
     'MISC': '#fea',
 }
 
+
 @st.cache(allow_output_mutation=True, show_spinner=False)
-def get_pipe():
-    model_name = "dslim/bert-base-NER"
+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")
+    pipe = transformers.pipeline("token-classification",
+                                 model=model,
+                                 tokenizer=tokenizer,
+                                 aggregation_strategy="simple")
     return pipe
 
+
 def parse_text(text, prediction):
     start = 0
     parsed_text = []
@@ -27,24 +31,25 @@ def parse_text(text, prediction):
     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.")
 
-default_text = "My name is John Smith. I work at Microsoft. I live in Paris. My favorite painting is the Mona Lisa."
+option = st.selectbox('Model', ("dslim/bert-base-NER", "flair/ner-english-fast", "Jean-Baptiste/camembert-ner"))
+st.write('Selected model:', option)
 
+default_text = "Xbox v PlayStation: Giants clash over Call of Duty: Xbox owner Microsoft has hit back at claims its plan to buy the maker of Call of Duty may unfairly affect its rivals, including Sony, which owns PlayStation."
 text = st.text_area('Enter text here:', value=default_text)
 submit = st.button('Predict')
 
 with st.spinner("Loading model..."):
-    pipe = get_pipe()
+    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)
-