import gradio as gr

from transformers import AutoTokenizer, AutoModelForCausalLM
import numpy as np
from transformers_cfg.grammar_utils import IncrementalGrammarConstraint
from transformers_cfg.generation.logits_process import GrammarConstrainedLogitsProcessor

MODEL_NAME = "gpt2"

if __name__ == "__main__":
    # Define your model and your tokenizer
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
    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.1, "p >= 10%"),
        (0.01, "p >= 1%"),
        (1e-20, "p < 1%"),
    ]

    label_to_color = {
        "p >= 10%": "green",
        "p >= 1%": "yellow",
        "p < 1%": "red"
    }


    def get_tokens_and_labels(prompt):
        """
        Given the prompt (text), return a list of tuples (decoded_token, label)
        """
        inputs = tokenizer([prompt], return_tensors="pt")

        # Load json grammar and create a GrammarConstrainedLogitsProcessor for each call
        with open("json_minimal.ebnf", "r") as file:
            grammar_str = file.read()
        grammar = IncrementalGrammarConstraint(grammar_str, "root", tokenizer)
        grammar_processor = GrammarConstrainedLogitsProcessor(grammar)

        outputs = model.generate(
            **inputs, max_new_tokens=50, repetition_penalty=1.05, return_dict_in_generate=True, output_scores=True, logits_processor=[grammar_processor]
        )
        # 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_tokens = outputs.sequences[:, input_length:]

        # Initialize the highlighted output with the prompt, which will have no color label
        highlighted_out = [(tokenizer.decode(token), None) for token in inputs.input_ids]
        # Get the (decoded_token, label) pairs for the generated tokens
        for token, proba in zip(generated_tokens[0], 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((tokenizer.decode(token), this_label))

        return highlighted_out


    demo = gr.Blocks()
    with demo:
        gr.Markdown(
            """
            # Grammar-Constrained Decoding with GPT-2 
            This is a demo of how you can constrain the output of a GPT-2 model using a formal grammar.
            Here we use a simple JSON grammar to constrain the output of the model to be valid JSON strings.
            The grammar is defined in `json_minimal.ebnf` and is written in the Extended Backus-Naur Form (EBNF).
            Internally, it relies on the library [`transformers-cfg`](https://github.com/epfl-dlab/transformers-CFG).
            For demo purpose, gpt2 is used, but you can use much larger models for better performance.
            """
        )

        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(label="Prompt", lines=3, value="This is a valid json string for http request:")
                button = gr.Button(f"Generate with json object using {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()