Spaces:
Paused
Paused
File size: 1,830 Bytes
29b560e 3b47068 a22e0d4 3b47068 29b560e a22e0d4 29b560e 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 |
import gradio as gr
import os
from huggingface_hub import login
from dotenv import load_dotenv
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
import torch
from threading import Thread
load_dotenv()
hf_token = os.getenv('HF_TOKEN', 'your-key-if-not-using-env')
login(hf_token, add_to_git_credential=True)
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)}
]
quant_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_quant_type="nf4"
)
tokenizer = AutoTokenizer.from_pretrained(MODEL)
tokenizer.pad_token = tokenizer.eos_token
streamer = TextIteratorStreamer(tokenizer)
model = AutoModelForCausalLM.from_pretrained(MODEL, device_map="auto", quantization_config=quant_config)
cplusplus = None
def translate(python):
inputs = tokenizer.apply_chat_template(messages_for(python), return_tensors="pt").to("cuda")
generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=80)
thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
cplusplus = ""
for chunk in streamer:
cplusplus += chunk
yield cplusplus
del inputs
torch.cuda.empty_cache()
demo = gr.Interface(fn=translate, inputs="code", outputs="markdown")
demo.launch()
|