Mike
add aggregation step
babd965
raw
history blame
No virus
1.17 kB
import gradio as gr
from transformers import pipeline
token_skill_classifier = pipeline(model="jjzha/jobbert_skill_extraction", aggregation_strategy="simple")
token_knowledge_classifier = pipeline(model="jjzha/jobbert_knowledge_extraction", aggregation_strategy="simple")
examples = [
"Knowing Python is a plus.",
]
def ner(text):
output_skills = token_skill_classifier(text)
for result in output_skills:
if result.get("entity_group"):
tag = result["entity_group"]
result["entity"] = tag + "-Skill"
del result["entity_group"]
output_knowledge = token_knowledge_classifier(text)
for result in output_knowledge:
if result.get("entity_group"):
tag = result["entity_group"]
result["entity"] = tag + "-Knowledge"
del result["entity_group"]
output = output_skills + output_knowledge
return {"text": text, "entities": output}
demo = gr.Interface(fn=ner,
inputs=gr.Textbox(placeholder="Enter sentence here..."),
outputs=gr.HighlightedText(),
examples=examples)
demo.launch()