import gradio as gr
from huggingface_hub import login
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import torch

MODEL = "m-a-p/OpenCodeInterpreter-DS-33B"
"bos_token": "<|begin_of_text|>",
CHAT_TEMPLATE = "{{ bos_token }}{% for message in messages %}{% 
if message['role'] == 'user' %}{{ '<|start_header_id|>user<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}{% elif message['role'] == 'assistant' %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}{% else %}{{ '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}"

system_message = "You are a computer programmer that can translate python code to C++ in order to improve performance"

def user_prompt_for(python):
    return f"Rewrite this python code to C++. You must search for the maximum performance. \
    Format your response in Markdown. This is the Code: \
    \n\n\
    {python}"

def messages_for(python):
    return [
        {"role": "system", "content": system_message},
        {"role": "user", "content": user_prompt_for(python)}
    ]

def apply_chat_template(messages):
    bos_token = "<｜begin▁of▁sentence｜>"
    result = bos_token
    for message in messages:
        if message['role'] == 'user':
            result += f"<|start_header_id|>user<|end_header_id|>\n\n{message['content']}<|eot_id|>"
        elif message['role'] == 'assistant':
            result += f"<|start_header_id|>assistant<|end_header_id|>\n\n{message['content']}<|eot_id|>"
        else:
            result += f"<|start_header_id|>{message['role']}<|end_header_id|>\n\n{message['content']}<|eot_id|>"
    return result


tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16, device_map="auto")
model.eval()
streamer = TextIteratorStreamer(tokenizer)

cplusplus = None
def translate(python):
    inputs = tokenizer(apply_chat_template(messages_for(python)), return_tensors="pt").to(model.device)
    generation_kwargs = dict(
                            inputs,
                            streamer=streamer,
                            max_new_tokens=1024,
                            do_sample=False,
                            pad_token_id=tokenizer.eos_token_id,
                            eos_token_id=tokenizer.eos_token_id
                            )
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    cplusplus = ""
    for chunk in streamer:
        cplusplus += chunk
        yield cplusplus

demo = gr.Interface(fn=translate, inputs="code", outputs="markdown")
demo.launch()