from transformers import pipeline
import gradio as gr

get_completion = pipeline("ner", model="dslim/bert-base-NER")

def ner(input):
    output = get_completion(input)
    return {"text": input, "entities": output}


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 ner(input):
    output = get_completion(input)#, parameters=None, ENDPOINT_URL=API_URL
    merged_tokens = merge_tokens(output)
    return {"text": input, "entities": merged_tokens}



if __name__ == "__main__":
    # Create the Gradio interface
    demo = gr.Interface(fn=ner,
                    inputs=[gr.Textbox(label="Text to find entities", lines=2)],
                    outputs=[gr.HighlightedText(label="Text with entities")],
                    title="Sanofi Digital M&S Demo: Name Entity Recognition",
                    description="Find entities using the `dslim/bert-base-NER` model under the hood!",
                    allow_flagging="never",
                    examples=["My name is Yunke, I lead a data science team at Sanofi. I live in Cambridge, Massachusattes", "My address is: 450 water street, Cambridge, MA",
                    "Barbie is a 2023 American fantasy comedy film directed by Greta Gerwig and written by Gerwig and Noah Baumbach.[6] Based on the Barbie fashion dolls by Mattel, it is the first live-action Barbie film after numerous computer-animated direct-to-video and streaming television films. The film follows Barbie (Margot Robbie) and Ken (Ryan Gosling) on a journey of self-discovery following an existential crisis. It features an ensemble cast that also includes America Ferrera, Kate McKinnon, Issa Rae, Rhea Perlman, and Will Ferrell.A live-action Barbie film was announced in September 2009 by Universal Pictures with Laurence Mark producing. Development began in April 2014, when Sony Pictures acquired the film rights. Following multiple writer and director changes and the casting of Amy Schumer and later Anne Hathaway as Barbie, the rights were transferred to Warner Bros. Pictures in October 2018. Robbie was cast in 2019, and Gerwig was announced as director and co-writer with Baumbach in 2021. The rest of the cast were announced in early 2022. Filming took place primarily at Warner Bros. Studios, Leavesden, in England and on the Venice Beach Skatepark in Los Angeles from March to July 2022.Barbie premiered at the Shrine Auditorium in Los Angeles on July 9, 2023, and was theatrically released in the United States on July 21, 2023, by Warner Bros. Pictures. Its simultaneous release with Christopher Nolan's Oppenheimer led to the 'Barbenheimer' phenomenon on social media, which encouraged audiences to see both films as a double feature. The film received mostly positive reviews, and has grossed over $495 million worldwide, becoming the ninth-highest grossing film of 2023"])
    # Launch the interface
    demo.launch() #share=True