Python本地运行DeepSeek-R1-Distill-Qwen-1.5B模型

1. 前言

这篇博客中,我将介绍如何使用 HuggingFace 的 transformers 库在本地运行 DeepSeek 1.5B 模型,并且处理一些常见的问题,比如缺少 <think> token 的情况。

2. 环境准备

首先,确保你已经安装了 transformerstorch 库:

pip install transformers torch

3. 加载模型与 Tokenizer

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

#加载tokenizer和模型
tokenizer = AutoTokenizer.from_pretrained("deepseek-aiDeepSeek-R1-Distill-Qwen-1.5B")
model = AutoModelForCausalLM.from_pretrained("deepseek-aiDeepSeek-R1-Distill-Qwen-1.5B")

4. 处理 <think> token

如果模型没有 <think> token,可以手动添加它,并且增加其嵌入:

#检查是否有 <think> token
if '<think>' not in tokenizer.get_vocab():
    tokenizer.add_tokens(['<think>'])#手动添加

#增加 <think> token 的嵌入
model.resize_token_embeddings(len(tokenizer))

5. 对话生成

你可以和模型进行对话,模型生成的输出将会是完整的:

#对话循环:与用户互动5次
for _ in range(5):
    user_input = input("YOU: ")

    #对用户输入进行tokenize
    inputs = tokenizer(user_input, return_tensors="pt", padding=True, truncation=True)

    #创建attention_mask
    attention_mask = inputs['attention_mask'] if 'attention_mask' in inputs else None

    #生成模型的回复
    with torch.no_grad():  # 关闭梯度计算,节省内存
        outputs = model.generate(
            inputs['input_ids'],
            max_length=150,  # 最大生成长度
            num_return_sequences=1,  # 返回一个生成结果
            pad_token_id=tokenizer.eos_token_id,  # 使用模型的结束token
            no_repeat_ngram_size=2,  # 防止生成重复的n-gram
            top_p=0.92,  # 使用nucleus sampling
            top_k=50,  # top-k采样
            temperature=0.7,  # 控制生成的多样性
            eos_token_id=tokenizer.eos_token_id,  # 确保模型生成到结束
            bos_token_id=tokenizer.bos_token_id,  # 确保模型从开始token生成
            attention_mask=attention_mask  # 使用正确的attention_mask
        )

    #解码输出并打印模型回复
    model_reply = tokenizer.decode(outputs[0], skip_special_tokens=True)

    print(f"Model: {model_reply}")

6. 总结

我介绍了如何通过 transformers 库加载和运行 DeepSeek 1.5B 模型,同时解决了 <think> token 缺失和填充问题。你可以在本地顺利运行此模型进行对话生成。

完整代码

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("deepseek-aiDeepSeek-R1-Distill-Qwen-1.5B")
model = AutoModelForCausalLM.from_pretrained("deepseek-aiDeepSeek-R1-Distill-Qwen-1.5B")

if '<think>' not in tokenizer.get_vocab():
    tokenizer.add_tokens(['<think>'])

model.eval()

model.resize_token_embeddings(len(tokenizer))

for _ in range(5):
    user_input = input("YOU: ")
    inputs = tokenizer(user_input, return_tensors="pt", padding=True, truncation=True)
    attention_mask = inputs['attention_mask'] if 'attention_mask' in inputs else None

    with torch.no_grad():
        outputs = model.generate(
            inputs['input_ids'],
            max_length=150,
            num_return_sequences=1,
            pad_token_id=tokenizer.eos_token_id,
            no_repeat_ngram_size=2,
            top_p=0.92,
            top_k=50,
            temperature=0.7,
            eos_token_id=tokenizer.eos_token_id,
            bos_token_id=tokenizer.bos_token_id,
            attention_mask=attention_mask
        )
    
    model_reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
    print(f"Model: {model_reply}")