Spaces:
Running on Zero
Running on Zero
File size: 2,464 Bytes
7119435 | 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 | """Hugging Face Space - Coding Assistant PRO - OpenCoder-8B-Instruct"""
import spaces
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
from threading import Thread
import traceback
model_path = "infly/OpenCoder-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True, torch_dtype=torch.bfloat16)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
class StopOnTokens(StoppingCriteria):
def __call__(self, input_ids, scores, **kwargs):
stop_ids = [96539]
for stop_id in stop_ids:
if input_ids[0][-1] == stop_id:
return True
return False
user_role = "user"
assistant_role = "assistant"
sft_end_token = "<|im_end|>"
@spaces.GPU()
def predict(message, history):
try:
stop = StopOnTokens()
model_messages = []
for i, item in enumerate(history):
model_messages.append({"role": user_role, "content": item[0]})
model_messages.append({"role": assistant_role, "content": item[1]})
model_messages.append({"role": user_role, "content": message})
model_inputs = tokenizer.apply_chat_template(model_messages, add_generation_prompt=True, return_tensors="pt").to(device)
streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
generate_kwargs = dict(input_ids=model_inputs, streamer=streamer, max_new_tokens=1024, do_sample=False, stopping_criteria=StoppingCriteriaList([stop]))
t = Thread(target=model.generate, kwargs=generate_kwargs)
t.start()
partial_message = ""
for new_token in streamer:
partial_message += new_token
if sft_end_token in partial_message:
break
yield partial_message
except Exception as e:
print(traceback.format_exc())
prompt_examples = [
"Write a quick sort algorithm in python.",
"Write a greedy snake game using pygame.",
"How to use numpy?"
]
chatbot = gr.Chatbot(label="Coding Assistant")
with gr.Blocks(theme=gr.themes.Soft(), fill_height=True) as demo:
gr.ChatInterface(predict, chatbot=chatbot, fill_height=True, examples=prompt_examples, cache_examples=False)
demo.launch()
|