Spaces:
Paused
Paused
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() | |