import torch
import gradio as gr
import re
import transformers
import peft
import traceback

from queue import Queue
from threading import Thread
import gc

CUDA_AVAILABLE = torch.cuda.is_available()

device = torch.device("cuda" if CUDA_AVAILABLE else "cpu")

tokenizer = transformers.AutoTokenizer.from_pretrained("cerebras/Cerebras-GPT-2.7B")
tokenizer.pad_token_id = 0

model = transformers.AutoModelForCausalLM.from_pretrained(
    "cerebras/Cerebras-GPT-2.7B", 
    load_in_8bit=True, 
    torch_dtype=torch.float16,
    device_map={'':0} if CUDA_AVAILABLE else 'auto',
)

model = peft.PeftModel.from_pretrained(
    model,
    'lxe/lora-cerebras-gpt2.7b-alpaca-shortprompt',
    torch_dtype=torch.float16
)

model.half()

# Streaming functionality taken from https://github.com/oobabooga/text-generation-webui/blob/master/modules/text_generation.py#L105

class Stream(transformers.StoppingCriteria):
    def __init__(self, callback_func=None):
        self.callback_func = callback_func

    def __call__(self, input_ids, scores) -> bool:
        if self.callback_func is not None:
            self.callback_func(input_ids[0])
        return False

class Iteratorize:
    """
    Transforms a function that takes a callback
    into a lazy iterator (generator).
    """
    def __init__(self, func, kwargs={}, callback=None):
        self.mfunc=func
        self.c_callback=callback
        self.q = Queue()
        self.sentinel = object()
        self.kwargs = kwargs
        self.stop_now = False

        def _callback(val):
            if self.stop_now:
                raise ValueError
            self.q.put(val)

        def gentask():
            try:
                ret = self.mfunc(callback=_callback, **self.kwargs)
            except ValueError:
                traceback.print_exc()
                pass
            except:
                traceback.print_exc()
                pass

            clear_torch_cache()
            self.q.put(self.sentinel)
            if self.c_callback:
                self.c_callback(ret)

        self.thread = Thread(target=gentask)
        self.thread.start()

    def __iter__(self):
        return self

    def __next__(self):
        obj = self.q.get(True,None)
        if obj is self.sentinel:
            raise StopIteration
        else:
            return obj

    def __del__(self):
        clear_torch_cache()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.stop_now = True
        clear_torch_cache()

def clear_torch_cache():
    gc.collect()
    if CUDA_AVAILABLE:
        torch.cuda.empty_cache()

def generate_text(
    history,  
    max_new_tokens, 
    do_sample, 
    temperature, 
    top_p, 
    top_k, 
    repetition_penalty, 
    typical_p, 
    num_beams
):
    # Create a conversation context of the last 4 entries in the history
    inp = ''.join([
        f"Human: {h[0]}\n\nAssistant: {'' if h[1] is None else h[1]}\n\n" for h in history[-4:]
    ]).strip()
     
    input_ids = tokenizer.encode(
        inp, 
        return_tensors='pt', 
        truncation=True, 
        add_special_tokens=False
    ).to(device) # type: ignore

    generate_params = {
        "input_ids": input_ids,
        "max_new_tokens": max_new_tokens,
        "do_sample": do_sample,
        "temperature": temperature,
        "top_p": top_p,
        "top_k": top_k,
        "repetition_penalty": repetition_penalty,
        "typical_p": typical_p,
        "num_beams": num_beams,
        "stopping_criteria": transformers.StoppingCriteriaList(),
        "pad_token_id": tokenizer.pad_token_id,
    }

    def generate_with_callback(callback=None, **kwargs):
        kwargs['stopping_criteria'].append(Stream(callback_func=callback))
        clear_torch_cache()
        with torch.no_grad():
            model.generate(**kwargs) # type: ignore

    def generate_with_streaming(**kwargs):
        return Iteratorize(generate_with_callback, kwargs, callback=None)

    with generate_with_streaming(**generate_params) as generator:
        for output in generator:
            new_tokens = len(output) - len(input_ids[0])
            reply = tokenizer.decode(output[-new_tokens:], skip_special_tokens=True)

            # If reply contains '^Human:' or '^Assistant:' 
            # then we have reached the end of the assistant's response
            stop_re = re.compile(r'^(Human|Assistant):', re.MULTILINE)
            if re.search(stop_re, reply):
                reply = ''.join(reply.split('\n')[:-1])
                history[-1][1] = reply.strip()
                yield history
                break

            # if reply contains 'EOS' then we have reached the end of the conversation
            if output[-1] in [tokenizer.eos_token_id]:
                yield history
                break

            history[-1][1] = reply.strip()
            yield history

with gr.Blocks() as demo:
    gr.Markdown("""
    ## 🐺🦙 Cerebras GPT-2.7B Alpcaca-Shortprompt LoRA Chatbot
    This is a very fast and relatively coherent (but hallucinating) chatbot. 
    It uses the [Cerebras-GPT-2.7B](https://huggingface.co/cerebras/Cerebras-GPT-2.7B), with a LoRA finetuned on the [Alpcaca Dataset](https://github.com/tloen/alpaca-lora/blob/main/alpaca_data_cleaned.json) dataset using a shorter prompt. 
    The chatbot keeps a very short conversation context of 4 entries. It's the fastest chatbot in the west! 
    More info [here](https://github.com/lxe/cerebras-lora-alpaca)
    """)
    with gr.Row():
        with gr.Column():
            chatbot = gr.Chatbot()
            msg = gr.Textbox(value="How old is the Earth?", placeholder="Type a message...")
            with gr.Row():
                clear = gr.Button("Clear")

        with gr.Column():
            max_new_tokens = gr.Slider(0, 2048, 200, step=1, label="max_new_tokens")
            do_sample = gr.Checkbox(True, label="do_sample")
            with gr.Row():
                with gr.Column():
                    temperature = gr.Slider(0, 2, 0.1, step=0.01, label="temperature")
                    top_p = gr.Slider(0, 1, 0.8, step=0.01, label="top_p")
                    top_k = gr.Slider(0, 100, 35, step=1, label="top_k")
                with gr.Column():
                    repetition_penalty = gr.Slider(0, 10, 1.1, step=0.01, label="repetition_penalty")
                    typical_p = gr.Slider(0, 1, 1, step=0.01, label="typical_p")
                    num_beams = gr.Slider(0, 10, 1, step=1, label="num_beams")

    def user(user_message, history):
        return "", history + [[user_message, None]]
    
    def fix_history(history):
        update_history = False
        for i, (user, bot) in enumerate(history):
            if bot is None:
                update_history = True
                history[i][1] = "_silence_"
        if update_history:
            chatbot.update(history) 

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        generate_text, inputs=[
            chatbot,
            max_new_tokens, 
            do_sample, 
            temperature, 
            top_p, 
            top_k, 
            repetition_penalty, 
            typical_p, 
            num_beams
        ], outputs=[chatbot],
    ).then(fix_history, chatbot)

    clear.click(lambda: None, None, chatbot, queue=False)

demo.queue().launch()