ysharma/stable-code-instruct-3b

1import argparse2import os3import spaces4 5 6import gradio as gr7 8import json9from threading import Thread10import torch11from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer12 13MAX_LENGTH = 409614DEFAULT_MAX_NEW_TOKENS = 102415 16 17def parse_args():18    parser = argparse.ArgumentParser()19    parser.add_argument("--base_model", type=str)  # model path20    parser.add_argument("--n_gpus", type=int, default=1)  # n_gpu21    return parser.parse_args()22 23@spaces.GPU()24def predict(message, history, system_prompt, temperature, max_tokens):25    global model, tokenizer, device26    instruction = "<|im_start|>system\nA chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.\n<|im_end|>\n"27    for human, assistant in history:28        instruction += '<|im_start|>user\n' + human + '\n<|im_end|>\n<|im_start|>assistant\n' + assistant29    instruction += '\n<|im_start|>user\n' + message + '\n<|im_end|>\n<|im_start|>assistant\n'30    problem = [instruction]31    stop_tokens = ["<|endoftext|>", "<|im_end|>"]32    streamer = TextIteratorStreamer(tokenizer, timeout=100.0, skip_prompt=True, skip_special_tokens=True)33    enc = tokenizer(problem, return_tensors="pt", padding=True, truncation=True)34    input_ids = enc.input_ids35    attention_mask = enc.attention_mask36 37    if input_ids.shape[1] > MAX_LENGTH:38        input_ids = input_ids[:, -MAX_LENGTH:]39 40    input_ids = input_ids.to(device)41    attention_mask = attention_mask.to(device)42    generate_kwargs = dict(43        {"input_ids": input_ids, "attention_mask": attention_mask},44        streamer=streamer,45        do_sample=True,46        top_p=0.95,47        temperature=0.5,48        max_new_tokens=DEFAULT_MAX_NEW_TOKENS,49    )50    t = Thread(target=model.generate, kwargs=generate_kwargs)51    t.start()52    outputs = []53    for text in streamer:54        outputs.append(text)55        if text in stop_tokens:56            break57        print(text)58        yield "".join(outputs)59 60 61 62if __name__ == "__main__":63    args = parse_args()64    tokenizer = AutoTokenizer.from_pretrained("stabilityai/stable-code-instruct-3b")65    model = AutoModelForCausalLM.from_pretrained("stabilityai/stable-code-instruct-3b", torch_dtype=torch.bfloat16)66    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')67    model = model.to(device)68    gr.ChatInterface(69        predict,70        title="Stable Code Instruct Chat - Demo",71        description="Chat Model Stable Code 3B",72        theme="soft",73        chatbot=gr.Chatbot(label="Chat History",),74        textbox=gr.Textbox(placeholder="input", container=False, scale=7),75        retry_btn=None,76        undo_btn="Delete Previous",77        clear_btn="Clear",78        additional_inputs=[79            gr.Textbox("A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.", label="System Prompt"),80            gr.Slider(0, 1, 0.9, label="Temperature"),81            gr.Slider(100, 2048, 1024, label="Max Tokens"),82        ],83        additional_inputs_accordion_name="Parameters",84    ).queue().launch()85