|
from transformers import AutoTokenizer, AutoModelForTokenClassification |
|
from transformers import pipeline |
|
tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER") |
|
model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER") |
|
|
|
nerp = pipeline("ner", model=model, tokenizer=tokenizer) |
|
|
|
import gradio as gr |
|
|
|
from helper import merge_tokens |
|
|
|
def ner(input): |
|
output = nerp(input) |
|
merged_tokens = merge_tokens(output) |
|
return {"text": input, "entities": merged_tokens} |
|
|
|
nerapp = gr.Interface(fn=ner, |
|
inputs=[gr.Textbox(label="Text to find entities", lines=2)], |
|
outputs=[gr.HighlightedText(label="Text with entities")], |
|
title="NER with dslim/bert-base-NER", |
|
description="Find entities using the `dslim/bert-base-NER` ", |
|
allow_flagging="never", |
|
examples=["My name is Sam, I'm building AI Apps and I live in Chennai", "My name is Lilly, I live in Chennai and work at Levitate"]) |
|
|
|
nerapp.launch() |
|
|