Pq234/deepseek-coder-33b-instruct

1import os2 3from threading import Thread4from typing import Iterator5 6import gradio as gr7import spaces8import torch9from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer10 11MAX_MAX_NEW_TOKENS = 204812DEFAULT_MAX_NEW_TOKENS = 102413total_count=014MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))15 16DESCRIPTION = """\17# DeepSeek-33B-Chat18 19This space demonstrates model [DeepSeek-Coder](https://huggingface.co/deepseek-ai/deepseek-coder-33b-instruct) by DeepSeek, a code model with 33B parameters fine-tuned for chat instructions.20 21**You can also try our 33B model in [official homepage](https://coder.deepseek.com/chat).**22"""23 24if not torch.cuda.is_available():25    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"26 27 28if torch.cuda.is_available():29    model_id = "deepseek-ai/deepseek-coder-33b-instruct"30    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")31    tokenizer = AutoTokenizer.from_pretrained(model_id)32    tokenizer.use_default_system_prompt = False33    34 35 36@spaces.GPU37def generate(38    message: str,39    chat_history: list[tuple[str, str]],40    system_prompt: str,41    max_new_tokens: int = 1024,42    temperature: float = 0.6,43    top_p: float = 0.9,44    top_k: int = 50,45    repetition_penalty: float = 1,46) -> Iterator[str]:47    global total_count48    total_count += 149    print(total_count)50    os.system("nvidia-smi")51    conversation = []52    if system_prompt:53        conversation.append({"role": "system", "content": system_prompt})54    for user, assistant in chat_history:55        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])56    conversation.append({"role": "user", "content": message})57 58    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")59    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:60        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]61        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")62    input_ids = input_ids.to(model.device)63 64    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)65    generate_kwargs = dict(66        {"input_ids": input_ids},67        streamer=streamer,68        max_new_tokens=max_new_tokens,69        do_sample=False,70        top_p=top_p,71        top_k=top_k,72        num_beams=1,73        # temperature=temperature,74        repetition_penalty=repetition_penalty,75        eos_token_id=3202176    )77    t = Thread(target=model.generate, kwargs=generate_kwargs)78    t.start()79 80    outputs = []81    for text in streamer:82        outputs.append(text)83        yield "".join(outputs).replace("<|EOT|>","")84 85 86chat_interface = gr.ChatInterface(87    fn=generate,88    additional_inputs=[89        gr.Textbox(label="System prompt", lines=6),90        gr.Slider(91            label="Max new tokens",92            minimum=1,93            maximum=MAX_MAX_NEW_TOKENS,94            step=1,95            value=DEFAULT_MAX_NEW_TOKENS,96        ),97        # gr.Slider(98        #     label="Temperature",99        #     minimum=0,100        #     maximum=4.0,101        #     step=0.1,102        #     value=0,103        # ),104        gr.Slider(105            label="Top-p (nucleus sampling)",106            minimum=0.05,107            maximum=1.0,108            step=0.05,109            value=0.9,110        ),111        gr.Slider(112            label="Top-k",113            minimum=1,114            maximum=1000,115            step=1,116            value=50,117        ),118        gr.Slider(119            label="Repetition penalty",120            minimum=1.0,121            maximum=2.0,122            step=0.05,123            value=1,124        ),125    ],126    stop_btn=gr.Button("Stop"),127    examples=[128        ["implement snake game using pygame"],129        ["Can you explain briefly to me what is the Python programming language?"],130        ["write a program to find the factorial of a number"],131    ],132)133 134with gr.Blocks(css="style.css") as demo:135    gr.Markdown(DESCRIPTION)136    chat_interface.render()137 138if __name__ == "__main__":139    demo.queue(max_size=20).launch()140