import gradio as gr
import torch
from threading import Thread
from typing import Iterator
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

MAX_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = 2048

# base_model_name = "m-a-p/OpenCodeInterpreter-DS-6.7B"
base_model_name = "m-a-p/OpenCodeInterpreter-DS-1.3B"
model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.float32, device_map="cpu", low_cpu_mem_usage=True)

tokenizer = AutoTokenizer.from_pretrained(base_model_name)

def format_prompt(message, history):
    system_prompt = "You are OpenCodeInterpreter, you are an expert programmer that helps to write code based on the user request, with concise explanations."
    prompt = []
    prompt.append({"role": "system", "content": system_prompt})
    for user_prompt, bot_response in history:
        prompt.extend([{"role": "user", "content": user_prompt}, {"role": "assistant", "content": bot_response}])
    prompt.append({"role": "user", "content": message})
    return prompt

def generate(prompt: str, history: list[tuple[str, str]], max_new_tokens: int = 1024, temperature: float = 0.3,
    top_p: float = 0.9, top_k: int = 50, repetition_penalty: float = 1 ) -> Iterator[str]:

    temperature = float(temperature)
    if temperature < 1e-2:
        temperature = 1e-2
    
    formatted_prompt = []
    formatted_prompt = format_prompt(prompt, history)
    
    input_ids = tokenizer.apply_chat_template(formatted_prompt, return_tensors="pt", add_generation_prompt=True)
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
    input_ids = input_ids.to(model.device)

    streamer = TextIteratorStreamer(tokenizer, timeout=15.0, skip_prompt=True, skip_special_tokens=True)
    generation_kwargs    = dict({"input_ids": input_ids}, streamer=streamer, max_new_tokens=max_new_tokens, do_sample=False, top_p=top_p, top_k=top_k, 
                              temperature=temperature, num_beams=1, repetition_penalty=repetition_penalty, eos_token_id=tokenizer.eos_token_id)

    t = Thread(target=model.generate, kwargs=generation_kwargs  )
    t.start()
   
    outputs = []
    for chunk in streamer:
        outputs.append(chunk)
        yield "".join(outputs).replace("<|EOT|>","")
        

mychatbot = gr.Chatbot(layout="bubble", avatar_images=["user.png", "botoci.png"], bubble_full_width=False, show_label=False, show_copy_button=True, likeable=True,)

additional_inputs =  additional_inputs=[
        gr.Slider(
            label="Max new tokens",
            minimum=1,
            maximum=MAX_MAX_NEW_TOKENS,
            step=1,
            value=512,
        ),
        gr.Slider(
            label="Temperature",
            minimum=0,
            maximum=1.0,
            step=0.1,
            value=0.3,
        ),
        gr.Slider(
            label="Top-p",
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=0.9,
        ),
        gr.Slider(
            label="Top-k",
            minimum=1,
            maximum=1000,
            step=1,
            value=50,
        ),
        gr.Slider(
            label="Repetition penalty",
            minimum=1.0,
            maximum=2.0,
            step=0.05,
            value=1,
        )]

iface  = gr.ChatInterface(fn=generate,
                        chatbot=mychatbot,
                        additional_inputs=additional_inputs,
                        description=" Running on CPU. The response may be slow for cpu environments. 🙏🏻",
                        retry_btn=None,
                        undo_btn=None
                       )

with gr.Blocks() as demo:
    gr.HTML("<center><h1>Tomoniai's Chat with OpenCodeInterpreter</h1></center>")
    iface.render()
    
demo.queue(max_size=10).launch(show_api=False)