browndw's picture
Update app.py
7a3bfd1
raw
history blame
2.56 kB
rom transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline
import gradio as gr
def ner_tagging(text):
model_name = "browndw/docusco-bert"
tokenizer = AutoTokenizer.from_pretrained(model_name, add_prefix_space=True)
model = AutoModelForTokenClassification.from_pretrained(model_name)
nlp = pipeline("ner", model=model, tokenizer=tokenizer)
output = []
text_2 = text.split(" ")
for i in range(len(text_2)):
ent = ner_results[i]["entity"]
if ent != "O":
output.extend([(text_2[i], ent), (" ", None)])
else:
output.extend([(text_2[i], None), (" ", None)])
return output
def get_entities(example):
model_name = "browndw/docusco-bert"
tokenizer = AutoTokenizer.from_pretrained(model_name, add_prefix_space=True)
model = AutoModelForTokenClassification.from_pretrained(model_name)
token_classifier = pipeline("token-classification", aggregation_strategy="simple", model=model, tokenizer=tokenizer)
output = []
i=0
prev_item = None
next_item = None
while i < (len(results)):
item = results[i]
p=i-1
n=i+1
if p > 0:
prev_item = results[p]
if n<(len(results)):
next_item = results[n]
if (i==0):
if item["start"]>0:
output.extend([(example[0:item["start"]], None)])
output.extend([(example[item["start"]:item["end"]], item["entity_group"])])
if (next_item!=None):
##verificar el tramo entre actual y siguiente
if(item["end"]!=next_item["start"]):
output.extend([(example[item["end"]:next_item["start"]], None)])
i=i+1
if item["end"] < len(example):
output.extend([(example[item["end"]:len(example)], None)])
return output
def greet(name):
return "Hello " + name + "!!"
iface = gr.Interface(fn=get_entities, inputs="text", outputs=['highlight'], examples=[['Jaws is a splendidly shrewd cinematic equation which not only gives you one or two very nasty turns when you least expect them but, possibly more important, knows when to make you think another is coming without actually providing it.'],
['Jaws is a splendidly shrewd cinematic equation which not only gives you one or two very nasty turns when you least expect them but, possibly more important, knows when to make you think another is coming without actually providing it.']], title="Test of docusco-bert ",)
iface.launch()