Spaces:
Paused
Paused
import gradio as gr | |
from huggingface_hub import login | |
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer | |
from threading import Thread | |
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) | |
streamer = TextIteratorStreamer(tokenizer) | |
cplusplus = None | |
def translate(python): | |
inputs = tokenizer(messages_for(python), return_tensors="pt") | |
generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=20) | |
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() | |