willdas/deepseek-coder-7b-instruct

1import os2from threading import Thread3from typing import Iterator4 5import gradio as gr6import spaces7import torch8from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer9 10MAX_MAX_NEW_TOKENS = 204811DEFAULT_MAX_NEW_TOKENS = 102412total_count=013MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))14 15DESCRIPTION = """\16# DeepSeek-6.7B-Chat17 18This space demonstrates model [DeepSeek-Coder](https://huggingface.co/deepseek-ai/deepseek-coder-6.7b-instruct) by DeepSeek, a code model with 6.7B parameters fine-tuned for chat instructions.19 20**You can also try our 33B model in [official homepage](https://coder.deepseek.com/chat).**21"""22 23if not torch.cuda.is_available():24    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"25 26 27if torch.cuda.is_available():28    model_id = "deepseek-ai/deepseek-coder-6.7b-instruct"29    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")30    tokenizer = AutoTokenizer.from_pretrained(model_id)31    tokenizer.use_default_system_prompt = False32    33 34 35@spaces.GPU36def generate(37    message: str,38    chat_history: list[tuple[str, str]],39    system_prompt: str,40    max_new_tokens: int = 1024,41    temperature: float = 0.6,42    top_p: float = 0.9,43    top_k: int = 50,44    repetition_penalty: float = 1,45) -> Iterator[str]:46    global total_count47    total_count += 148    print(total_count)49    if total_count % 50 == 0 :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=None,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=5).launch()140