Arturo Jiménez de los Galanes Reguillos
Refactor to fit model recomended usage
314c465
raw
history blame
1.69 kB
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"
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)}
]
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()