Weixin-Liang
add model weights
88734c0
raw
history blame
1.26 kB
# REF: https://gradio.app/named_entity_recognition/
from transformers import pipeline
import gradio as gr
model_name="xlm-roberta-base"
# model_name="roberta-large"
from transformers import AutoTokenizer, AutoModelForTokenClassification
label_list= ['literal',"metaphoric"]
label_dict_relations={ i : l for i, l in enumerate(label_list) }
PATH = "./saved-models/my_model"
model_metaphor_detection = AutoModelForTokenClassification.from_pretrained(PATH, id2label=label_dict_relations)
tokenizer = AutoTokenizer.from_pretrained(model_name)
pipeline_metaphors=pipeline("ner", model=model_metaphor_detection, tokenizer=tokenizer, aggregation_strategy="simple")
examples = [
"It would change the trajectory of your legal career.",
"Washington and the media just explodes on you, you just don’t know where you are at the moment",
"Those statements are deeply concerning.",
]
def ner(text):
output = pipeline_metaphors(text)
# change name
for x in output:
x['entity'] = x['entity_group']
return {"text": text, "entities": output}
demo = gr.Interface(ner,
gr.Textbox(placeholder="Enter sentence here..."),
gr.HighlightedText(),
examples=examples)
demo.launch(share=True)