|
import gradio as gr |
|
from transformers import pipeline |
|
|
|
get_completion = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6") |
|
get_ner = pipeline("ner", model="dslim/bert-base-NER") |
|
get_zero = pipeline( |
|
"zero-shot-classification", model="MoritzLaurer/DeBERTa-v3-base-mnli-fever-anli" |
|
) |
|
|
|
|
|
def summarize_text(input): |
|
output = get_completion(input) |
|
return output[0]["summary_text"] |
|
|
|
|
|
def merge_tokens(tokens): |
|
merged_tokens = [] |
|
for token in tokens: |
|
if ( |
|
merged_tokens |
|
and token["entity"].startswith("I-") |
|
and merged_tokens[-1]["entity"].endswith(token["entity"][2:]) |
|
): |
|
|
|
last_token = merged_tokens[-1] |
|
last_token["word"] += token["word"].replace("##", "") |
|
last_token["end"] = token["end"] |
|
last_token["score"] = (last_token["score"] + token["score"]) / 2 |
|
else: |
|
|
|
merged_tokens.append(token) |
|
return merged_tokens |
|
|
|
|
|
def named_entity_recognition(input): |
|
output = get_ner(input) |
|
merged_output = merge_tokens(output) |
|
return {"text": input, "entities": output} |
|
|
|
|
|
def zero_shot_pred(text, check_labels): |
|
output = get_zero(text, check_labels) |
|
return output |
|
|
|
|
|
def label_score_dict(text, check_labels): |
|
zero_shot_out = zero_shot_pred(text, check_labels) |
|
out = {} |
|
for i, j in zip(zero_shot_out["labels"], zero_shot_out["scores"]): |
|
out.update({i: j}) |
|
print(out) |
|
return out |
|
|
|
|
|
interface_summarise = gr.Interface( |
|
fn=summarize_text, |
|
inputs=[gr.Textbox(label="Text to summarise", lines=5)], |
|
outputs=[gr.Textbox(label="Summary")], |
|
title="Text Summarizer", |
|
description="Summary of text via `distillBART-CNN` model!", |
|
) |
|
|
|
interface_ner = gr.Interface( |
|
fn=named_entity_recognition, |
|
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` model under the hood!", |
|
allow_flagging="never", |
|
examples=[ |
|
"Tim Cook is the CEO of Apple, stays in California and makes iPhones ", |
|
"My name is Bose and I am a physicist living in Delhi", |
|
], |
|
) |
|
|
|
interface_zero_shot = gr.Interface( |
|
fn=label_score_dict, |
|
inputs=[ |
|
gr.Textbox(label="Text to classify", lines=2), |
|
gr.Textbox(label="Check for labels"), |
|
], |
|
outputs=gr.Label(num_top_classes=4), |
|
title="Zero-Shot Preds using DeBERTa-v3-base-mnli", |
|
description="Classify sentence on self defined target vars", |
|
examples=[ |
|
[ |
|
"Last week I upgraded my iOS version and ever since then my phone has been overheating whenever I use your app.", |
|
"mobile, website, billing, account access", |
|
], |
|
|
|
], |
|
) |
|
|
|
demo = gr.TabbedInterface( |
|
[interface_summarise, interface_ner, interface_zero_shot], |
|
["Text Summary ", "Named Entity Recognition", "Zero Shot Classifications"], |
|
) |
|
|
|
if __name__ == "__main__": |
|
demo.launch(enable_queue=True) |
|
|
