Spaces:
Running
Running
File size: 3,115 Bytes
9d250e2 859baec 9d250e2 859baec 9d250e2 859baec 9d250e2 859baec 9d250e2 859baec 9d250e2 859baec 9d250e2 859baec 9d250e2 859baec 9d250e2 859baec 9d250e2 859baec 9d250e2 859baec 9d250e2 859baec 9d250e2 |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import torch
import gradio as gr
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
MODEL_ID = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
MODEL_NAME = MODEL_ID.split("/")[-1]
CONTEXT_LENGTH = 16000
DESCRIPTION = f"This is {MODEL_NAME} model designed for testing thinking for general AI tasks. <br>当前仅提供 HuggingFace 版部署实例，有算力的可自行克隆至本地或复刻至购买了 GPU 环境的账号测试"
def predict(
message,
history,
system_prompt,
temperature,
max_new_tokens,
top_k,
repetition_penalty,
top_p,
):
# Format history with a given chat template
stop_tokens = ["<|endoftext|>", "<|im_end|>", "|im_end|"]
instruction = "<|im_start|>system\n" + system_prompt + "\n<|im_end|>\n"
for user, assistant in history:
instruction += f"<|im_start|>user\n{user}\n<|im_end|>\n<|im_start|>assistant\n{assistant}\n<|im_end|>\n"
instruction += f"<|im_start|>user\n{message}\n<|im_end|>\n<|im_start|>assistant\n"
print(instruction)
streamer = TextIteratorStreamer(
tokenizer,
skip_prompt=True,
skip_special_tokens=True,
)
enc = tokenizer(instruction, return_tensors="pt", padding=True, truncation=True)
input_ids, attention_mask = enc.input_ids, enc.attention_mask
if input_ids.shape[1] > CONTEXT_LENGTH:
input_ids = input_ids[:, -CONTEXT_LENGTH:]
attention_mask = attention_mask[:, -CONTEXT_LENGTH:]
generate_kwargs = dict(
input_ids=input_ids.to(device),
attention_mask=attention_mask.to(device),
streamer=streamer,
do_sample=True,
temperature=temperature,
max_new_tokens=max_new_tokens,
top_k=top_k,
repetition_penalty=repetition_penalty,
top_p=top_p,
)
t = Thread(target=model.generate, kwargs=generate_kwargs)
t.start()
outputs = []
for new_token in streamer:
outputs.append(new_token)
if new_token in stop_tokens:
break
yield "".join(outputs)
if __name__ == "__main__":
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto")
# Create Gradio interface
gr.ChatInterface(
predict,
title=f"{MODEL_NAME} 部署实例",
description=DESCRIPTION,
additional_inputs_accordion=gr.Accordion(label="⚙️ Parameters", open=False),
additional_inputs=[
gr.Textbox(
"You are a useful assistant. first recognize user request and then reply carfuly and thinking",
label="System prompt",
),
gr.Slider(0, 1, 0.6, label="Temperature"),
gr.Slider(0, 32000, 10000, label="Max new tokens"),
gr.Slider(1, 80, 40, label="Top K sampling"),
gr.Slider(0, 2, 1.1, label="Repetition penalty"),
gr.Slider(0, 1, 0.95, label="Top P sampling"),
],
).queue().launch()
|