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

MODEL = "NTQAI/Nxcode-CQ-7B-orpo"

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 python Code: \
    \n\n\
    {python}"

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

tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype="auto", device_map="auto")

decode_kwargs = dict(skip_special_tokens=True)
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, decode_kwargs=decode_kwargs)

cplusplus = None
def translate(python):
    inputs = tokenizer.apply_chat_template(
                        messages_for(python),
                        add_generation_prompt=True,
                        return_tensors="pt").to(model.device)

    generation_kwargs = dict(
        input_ids=inputs,
        streamer=streamer,
        max_new_tokens=512,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        num_return_sequences=1,
        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()