Arturo Jiménez de los Galanes Reguillos
Remove unused packages
6d8bd5f
raw
history blame
1.31 kB
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()