Update app.py
Browse files
app.py
CHANGED
@@ -1,5 +1,6 @@
|
|
1 |
|
2 |
import gradio as gr
|
|
|
3 |
from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer, AutoModelForTokenClassification
|
4 |
|
5 |
# Load your custom model and tokenizer
|
@@ -11,19 +12,35 @@ qa_tokenizer = AutoTokenizer.from_pretrained(qa_model_name)
|
|
11 |
|
12 |
token_classification_model = AutoModelForTokenClassification.from_pretrained(token_classification_model_name)
|
13 |
token_classification_tokenizer = AutoTokenizer.from_pretrained(token_classification_model_name)
|
14 |
-
|
15 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
16 |
def predict(pipeline_type, question, context):
|
17 |
if pipeline_type == "question-answering":
|
18 |
qa_pipeline = pipeline("question-answering", model=qa_model, tokenizer=qa_tokenizer)
|
19 |
result = qa_pipeline(question=question, context=context)
|
20 |
-
response = [(result['answer'], result
|
21 |
-
return response
|
22 |
elif pipeline_type == "token-classification":
|
23 |
-
token_classification_pipeline = pipeline("
|
24 |
result = token_classification_pipeline(context)
|
25 |
-
highlighted_text = {"text": context, "entities": result}
|
26 |
-
|
|
|
27 |
|
28 |
# Create a Gradio Interface with dropdown and two text inputs
|
29 |
iface = gr.Interface(
|
@@ -33,8 +50,11 @@ iface = gr.Interface(
|
|
33 |
"text",
|
34 |
"text"
|
35 |
],
|
36 |
-
outputs=gr.Highlight()
|
37 |
)
|
38 |
|
|
|
|
|
|
|
39 |
# Launch the interface
|
40 |
-
iface.launch()
|
|
|
1 |
|
2 |
import gradio as gr
|
3 |
+
import pandas as pd
|
4 |
from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer, AutoModelForTokenClassification
|
5 |
|
6 |
# Load your custom model and tokenizer
|
|
|
12 |
|
13 |
token_classification_model = AutoModelForTokenClassification.from_pretrained(token_classification_model_name)
|
14 |
token_classification_tokenizer = AutoTokenizer.from_pretrained(token_classification_model_name)
|
15 |
+
def tabulazier(output):
|
16 |
+
output_comb = []
|
17 |
+
for ind, entity in enumerate(output):
|
18 |
+
if ind == 0:
|
19 |
+
output_comb.append(entity)
|
20 |
+
elif output[ind]["start"] == output[ind-1]["end"] and output[ind]["entity_group"] == output[ind-1]["entity_group"]:
|
21 |
+
output_comb[-1]["word"] = output_comb[-1]["word"] + output[ind]["word"]
|
22 |
+
output_comb[-1]["end"] = output[ind]["end"]
|
23 |
+
else:
|
24 |
+
output_comb.append(entity)
|
25 |
+
|
26 |
+
df = pd.DataFrame(output_comb)
|
27 |
+
df['word'] = df['word'].str.replace('#', '')
|
28 |
+
return df
|
29 |
+
|
30 |
+
|
31 |
+
# Define a function for inference based on pipeline type
|
32 |
def predict(pipeline_type, question, context):
|
33 |
if pipeline_type == "question-answering":
|
34 |
qa_pipeline = pipeline("question-answering", model=qa_model, tokenizer=qa_tokenizer)
|
35 |
result = qa_pipeline(question=question, context=context)
|
36 |
+
response = [(result['answer'], result.get('score', None))]
|
37 |
+
return [response, response]
|
38 |
elif pipeline_type == "token-classification":
|
39 |
+
token_classification_pipeline = pipeline("ner", model=token_classification_model, tokenizer=token_classification_tokenizer, aggregation_strategy="simple")
|
40 |
result = token_classification_pipeline(context)
|
41 |
+
highlighted_text = {"text": context, "entities": result}
|
42 |
+
table=tabulazier(result)
|
43 |
+
return [gr.HighlightedText(highlighted_text), table]
|
44 |
|
45 |
# Create a Gradio Interface with dropdown and two text inputs
|
46 |
iface = gr.Interface(
|
|
|
50 |
"text",
|
51 |
"text"
|
52 |
],
|
53 |
+
outputs=[gr.Highlight(), gr.Dataframe()]
|
54 |
)
|
55 |
|
56 |
+
|
57 |
+
|
58 |
+
|
59 |
# Launch the interface
|
60 |
+
iface.launch(debug=False)
|