bert-pulc / app.py
Umesh's picture
Update app.py
6c0c4b3
Raw History Blame Contribute Delete
2.86 kB
import gradio as gr
from transformers import pipeline
import spacy
import spacy.cli
spacy.cli.download("en_core_web_lg")
nlp = spacy.load("en_core_web_lg")
model_checkpoint ="Umesh/bert-pulc"
token_classifier = pipeline("text-classification", model=model_checkpoint, aggregation_strategy="simple")
def clean_text(text):
'''
Cleans the input text
'''
text = re.sub(r'\\n', ' ', text)
text = text.encode('ascii', 'ignore').decode()
text = text.strip()
text=re.sub(r'http\S+', '', text)
text = re.sub("[^,.:1-9A-Za-z]+", " ", text)
text = re.sub(" +", " ", text)
return text
def main(everyline):
doc = nlp(everyline) #identifies named entities
#print('Entities', [(ent.text, ent.label_) for ent in doc.ents])
#print('Tokens', [(t.text, t.ent_type_, t.ent_iob) for t in doc])
#print("------------------------------")
finalString = ""
for t in doc:
# skipping inside labels
if int(t.ent_iob) == 1:
continue
# this takes care of the first piece of string in a text
# Also we identify the beginning label which is label 3 in the following if else statement.
if finalString == "":
if str(t.ent_type_) == "PERSON":
finalString = "<person>"
elif str(t.ent_type_) == 'GPE':
finalString = "<loc>"
elif str(t.ent_type_) == 'LOC':
finalString = "<loc>"
elif str(t.ent_type_) == 'NORP':
finalString = "<norp>"
elif str(t.ent_type_) == 'FAC':
finalString = "<fac>"
elif str(t.ent_type_) == 'ORG':
finalString = "<org>"
else:
finalString = str(t.text)
elif int(t.ent_iob) == 3:
if str(t.ent_type_) == "PERSON":
finalString += " <person>"
elif str(t.ent_type_) == 'GPE':
finalString += " <loc>"
elif str(t.ent_type_) == 'LOC':
finalString += " <loc>"
elif str(t.ent_type_) == 'NORP':
finalString += " <norp>"
elif str(t.ent_type_) == 'FAC':
finalString += " <fac>"
elif str(t.ent_type_) == 'ORG':
finalString += " <org>"
else:
finalString += " " + str(t.text)
else:
finalString += " " + str(t.text)
return finalString
def inference(text):
output = token_classifier(text)
return {p["label"]: p["score"] for p in output}
demo = gr.Interface(inference,
gr.Textbox(placeholder="Enter sentence here..."),
gr.HighlightedText(),
allow_flagging = False,
title="Text Classification-PULF")
demo.launch()