add torch dep, add model choices

app.py

@@ -1,23 +1,28 @@
-import gradio as gr
 from transformers import pipeline
+import gradio as gr
 
 
-ner_pipeline = pipeline("ner", model="dslim/bert-base-NER")
-
-examples = [
-    "Does Chicago have any stores and does Joe live here?",
-]
+# List of NER models
+models = ["dslim/bert-base-NER", "dslim/bert-base-NER-uncased", "dslim/bert-large-NER"]
 
 
-def ner(text):
+def ner(text, model_choice):
+    ner_pipeline = pipeline("ner", model=model_choice)
     output = ner_pipeline(text)
     return {"text": text, "entities": output}
 
 
+examples = [
+    "Does Chicago have any stores and does Joe live here?",
+]
+
 demo = gr.Interface(
-    ner,
-    gr.Textbox(placeholder="Enter sentence here..."),
-    gr.HighlightedText(),
+    fn=ner,
+    inputs=[
+        gr.Textbox(placeholder="Enter sentence here..."),
+        gr.Dropdown(choices=models, label="Choose NER Model"),
+    ],
+    outputs=gr.HighlightedText(),
     examples=examples,
 )
 

requirements.txt

@@ -1,2 +1,3 @@
 transformers
 gradio
+torch