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_caption = pipeline("image-to-text")
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}


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"
                    ])

demo = gr.TabbedInterface([
    interface_summarise, 
    interface_ner],
    ["Text Summary ",
     "Named Entity Recognition"
    ])

if __name__ == "__main__":
    demo.launch(enable_queue=True)