Spaces:
Paused
Paused
File size: 2,878 Bytes
29b560e 3b47068 314c465 a22e0d4 e93ccdc 3b47068 a22e0d4 fde3482 29b560e a22e0d4 fde3482 a22e0d4 314c465 a22e0d4 fde3482 0a193bf a22e0d4 29b560e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
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"
"bos_token": "<|begin_of_text|>",
CHAT_TEMPLATE = "{{ bos_token }}{% for message in messages %}{%
if message['role'] == 'user' %}{{ '<|start_header_id|>user<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}{% elif message['role'] == 'assistant' %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}{% else %}{{ '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}"
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)}
]
def apply_chat_template(messages):
bos_token = "<|begin▁of▁sentence|>"
result = bos_token
for message in messages:
if message['role'] == 'user':
result += f"<|start_header_id|>user<|end_header_id|>\n\n{message['content']}<|eot_id|>"
elif message['role'] == 'assistant':
result += f"<|start_header_id|>assistant<|end_header_id|>\n\n{message['content']}<|eot_id|>"
else:
result += f"<|start_header_id|>{message['role']}<|end_header_id|>\n\n{message['content']}<|eot_id|>"
return result
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()
|