import gradio as gr

from transformers import AutoTokenizer, AutoModelForCausalLM
import numpy as np


MODEL_NAME = "microsoft/Phi-3-mini-4k-instruct"

my_token = "h"+"f"+"_"+"yFicBqLnJDUkEIpccOIKpYMecxvPoTiUpG"


if __name__ == "__main__":
    # Define your model and your tokenizer
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, trust_remote_code=True)  # was AutoModelForSeq2SeqLM in case of "google/flan-t5-base"
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token_id = tokenizer.eos_token_id
        model.config.pad_token_id = model.config.eos_token_id

    # Define your color-coding labels; if prob > x, then label = y; Sorted in descending probability order!
    probs_to_label = [
        (0.9, "p >= 90%"),
        (0.8, "p >= 80%"),
        (0.7, "p >= 70%"),
        (0.6, "p >= 60%"),
        (0.5, "p >= 50%"),
        (0.4, "p >= 40%"),
        (0.3, "p >= 30%"),
        (0.2, "p >= 20%"),
        (0.1, "p >= 10%"),
        (0.0, "p >= 00%")
    ]

    label_to_color = {
        "p >= 90%": "#11d9d2",
        "p >= 80%": "#11b4d9",
        "p >= 80%": "#11d9a0",
        "p >= 70%": "#11d954",
        "p >= 60%": "#4dd911",
        "p >= 50%": "#a0d911",
        "p >= 40%": "#d5d911",
        "p >= 30%": "#d9c111",
        "p >= 20%": "#d99a11",
        "p >= 10%": "#d97211",
        "p >= 00%": "#d91111"
    }

    def get_tokens_and_labels(prompt):
        """
        Given the prompt (text), return a list of tuples (decoded_token, label)
        """
        inputs = tokenizer([prompt], return_tensors="pt")
        outputs = model.generate(
            **inputs, max_new_tokens=50, return_dict_in_generate=True, output_scores=True
        )
        # Important: don't forget to set `normalize_logits=True` to obtain normalized probabilities (i.e. sum(p) = 1)
        transition_scores = model.compute_transition_scores(outputs.sequences, outputs.scores, normalize_logits=True)
        transition_proba = np.exp(transition_scores)
        # We only have scores for the generated tokens, so pop out the prompt tokens
        input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1]
        generated_ids = outputs.sequences[:, input_length:]
        generated_tokens = tokenizer.convert_ids_to_tokens(generated_ids[0])

        # Important: you might need to find a tokenization character to replace (e.g. "Ġ" for BPE) and get the correct
        # spacing into the final output 👼
        if model.config.is_encoder_decoder:
            highlighted_out = []
        else:
            input_tokens = tokenizer.convert_ids_to_tokens(inputs.input_ids[0])
            highlighted_out = [(token.replace("▁", " "), None) for token in input_tokens]
        # Get the (decoded_token, label) pairs for the generated tokens
        for token, proba in zip(generated_tokens, transition_proba[0]):
            this_label = None
            assert 0. <= proba <= 1.0
            for min_proba, label in probs_to_label:
                if proba >= min_proba:
                    this_label = label
                    break
            highlighted_out.append((token.replace("▁", " "), this_label))

        return highlighted_out

    demo = gr.Blocks()
    with demo:
        gr.Markdown(
            """
            # 🌈 Color-Coded Text Generation 🌈

            This is a demo of how you can obtain the probabilities of each generated token, and use them to
            color code the model output. Internally, it relies on
            [`compute_transition_scores`](https://huggingface.co/docs/transformers/main/en/main_classes/text_generation#transformers.GenerationMixin.compute_transition_scores),
            which was added in `transformers` v4.26.0.

            ⚠️ For instance, with the pre-populated input and its color-coded output, you can see that
            `google/flan-t5-base` struggles with arithmetics.

            🤗 Feel free to clone this demo and modify it to your needs 🤗
            """
        )

        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="Prompt",
                    lines=3,
                    value=(
                        "Answer the following question by reasoning step-by-step. The cafeteria had 23 apples. "
                        "If they used 20 for lunch and bought 6 more, how many apples do they have?"
                    ),
                )
                button = gr.Button(f"Generate with {MODEL_NAME}")
            with gr.Column():
                highlighted_text = gr.HighlightedText(
                    label="Highlighted generation",
                    combine_adjacent=True,
                    show_legend=True,
                ).style(color_map=label_to_color)

        button.click(get_tokens_and_labels, inputs=prompt, outputs=highlighted_text)


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