dslim commited on
Commit
0a4290b
1 Parent(s): 2efee62

add torch dep, add model choices

Browse files
Files changed (2) hide show
  1. app.py +15 -10
  2. requirements.txt +1 -0
app.py CHANGED
@@ -1,23 +1,28 @@
1
- import gradio as gr
2
  from transformers import pipeline
 
3
 
4
 
5
- ner_pipeline = pipeline("ner", model="dslim/bert-base-NER")
6
-
7
- examples = [
8
- "Does Chicago have any stores and does Joe live here?",
9
- ]
10
 
11
 
12
- def ner(text):
 
13
  output = ner_pipeline(text)
14
  return {"text": text, "entities": output}
15
 
16
 
 
 
 
 
17
  demo = gr.Interface(
18
- ner,
19
- gr.Textbox(placeholder="Enter sentence here..."),
20
- gr.HighlightedText(),
 
 
 
21
  examples=examples,
22
  )
23
 
 
 
1
  from transformers import pipeline
2
+ import gradio as gr
3
 
4
 
5
+ # List of NER models
6
+ models = ["dslim/bert-base-NER", "dslim/bert-base-NER-uncased", "dslim/bert-large-NER"]
 
 
 
7
 
8
 
9
+ def ner(text, model_choice):
10
+ ner_pipeline = pipeline("ner", model=model_choice)
11
  output = ner_pipeline(text)
12
  return {"text": text, "entities": output}
13
 
14
 
15
+ examples = [
16
+ "Does Chicago have any stores and does Joe live here?",
17
+ ]
18
+
19
  demo = gr.Interface(
20
+ fn=ner,
21
+ inputs=[
22
+ gr.Textbox(placeholder="Enter sentence here..."),
23
+ gr.Dropdown(choices=models, label="Choose NER Model"),
24
+ ],
25
+ outputs=gr.HighlightedText(),
26
  examples=examples,
27
  )
28
 
requirements.txt CHANGED
@@ -1,2 +1,3 @@
1
  transformers
2
  gradio
 
 
1
  transformers
2
  gradio
3
+ torch