add torch dep, add model choices
Browse files- app.py +15 -10
- requirements.txt +1 -0
app.py
CHANGED
@@ -1,23 +1,28 @@
|
|
1 |
-
import gradio as gr
|
2 |
from transformers import pipeline
|
|
|
3 |
|
4 |
|
5 |
-
|
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 |
-
|
20 |
-
|
|
|
|
|
|
|
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
|