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:]):
            # If current token continues the entity of the last one, merge them
            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:
            # Otherwise, add the token to the list
            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"],
                        # "My name is Bose and I am a physicist living in Delhi"
                    ])

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)