RanjithkumarPanjabikesan's picture
Create app.py
ba710bc verified
raw
history blame
1.28 kB
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")
#Create Named Entity Recognition Pipeline
nerp = pipeline("ner", model=model, tokenizer=tokenizer)
#Build the Named Entity Recognition App
import gradio as gr
#Import Merging Tokens function from Helper to Display Output Relevant for User
from helper import merge_tokens
#Define Named Entity Recognition Function
def ner(input):
output = nerp(input)
merged_tokens = merge_tokens(output)
return {"text": input, "entities": merged_tokens}
#Set up the User Interface and Launch
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 Ranjith, I love AI and I live in Chennai", "My name is Akshay, I live in Germany and work at Cocacola"])
nerapp.launch()