import gradio as gr
import outlines
import transformers
import torch
from threading import Thread

pipe = transformers.pipeline("text-generation", "HuggingFaceTB/SmolLM-1.7B-Instruct", torch_dtype=torch.float32)
outlines_tokenizer = outlines.models.TransformerTokenizer(pipe.tokenizer)

def string_to_acrostic_grammar(s, dash_initial=True):
    # this will convert a string to a CFG grammar
    chars = filter(str.isalpha, s.upper())
    grammar_rules = [('"- " ' if dash_initial else '') + f'"{char}" /[^-\\r\\n]+/ "\\n"' for char in chars]
    return "?start: " + " ".join(grammar_rules)

def is_this_prompt_a_list(prompt):
    # this will check if the prompt is a list
    # ask the model if the prompt is a list, by constraining the generation to yes or no about a question whether the prompt is a list
    question = f'You are trying to understand the desired format of output for a prompt, whether it will be a list or a story. The prompt:\n```{prompt}```\n\nIs this prompt asking for short phrases in a list, or long sentences in a story?'
    grammar = '?start: ("list" | "story")'
    cfg_logits_processor = outlines.processors.CFGLogitsProcessor(grammar, outlines_tokenizer)
    output = pipe([{"role": "user", "content": question}, {"role": "assistant", "content": "The output to this prompt is a "}], logits_processor=transformers.LogitsProcessorList([cfg_logits_processor]), max_new_tokens=10,)
    response = output[0]['generated_text'][-1]['content'].split()[-1]
    # the last word is the answer
    print("is this prompt a list?", response)
    return response == "list"

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    acrostic,
    max_tokens,
    temperature,
    top_p,
):
    print({"message": message, "history": history, "system_message": system_message, "acrostic": acrostic, "max_tokens": max_tokens, "temperature": temperature, "top_p": top_p})
    # this will generate a response to the message
    prompt = f"<|im_start|>system\n{system_message}<|im_end|>\n<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    grammar = string_to_acrostic_grammar(acrostic, dash_initial=is_this_prompt_a_list(prompt))
    acrostic_logits_processor = outlines.processors.CFGLogitsProcessor(grammar, outlines_tokenizer)
    streamer = transformers.TextIteratorStreamer(pipe.tokenizer, skip_prompt=True, decode_kwargs={"skip_special_tokens": True})
    current_inputs = []
    # take the current inputs, and for every item in the history (which is a list of [x,y], add it to the current inputs like so: {"role": "user", "content": x), {"role": "assistant", "content": y}
    for x, y in history:
        current_inputs.append({"role": "user", "content": x})
        current_inputs.append({"role": "assistant", "content": y})
    # add the current inputs to the inputs
    inputs = current_inputs + [{"role": "user", "content": prompt}]

    generation_kwargs = dict(text_inputs=inputs, logits_processor=transformers.LogitsProcessorList([acrostic_logits_processor]), streamer=streamer, max_new_tokens=max_tokens, temperature=temperature, top_p=top_p, do_sample=True)
    thread = Thread(target=pipe, kwargs=generation_kwargs)
    thread.start()
    # this will generate a response to the message
    # TODO: figure out why skip special tokens doesn't skip special tokens
    special_tokens = set([str(v) for v in pipe.tokenizer.special_tokens_map.values()])
    response = ""
    for new_text in streamer:
        if new_text not in special_tokens:
            response += new_text
        yield response


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Textbox(value="I love you", label="acrostic"),
        gr.Slider(minimum=1, maximum=8192, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.2, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)


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