PythonでDeepSeekをローカルで実行するコード

DeepSeekがバズっている。会社でこれをローカルマシンで動かすように言われたので、どのように書けばいいか調べてまとめてみた。

※なお、まだコードは実際に実行して試していない。自分が持っているマシンじゃ到底動くはずもないので、会社で試す予定。ここにまとめているのは一旦自分が会社でカンニングするためのメモ用。

vLLM使う場合

vllmを使う場合は、vLLMサーバーを立ち上げ、そのサーバーに対してAPIリクエストをする形式で実行するみたい。

main.py

import subprocess
import time
import requests

def main():
    # 1) vLLMサーバーをバックグラウンド起動
    #    （Windowsや環境によってはshell=Trueが必要になる場合もあります）
    p = subprocess.Popen(["vllm", "serve", "deepseek-ai/DeepSeek-R1"])
    
    # 2) サーバーが立ち上がるまで少し待機
    time.sleep(10)  # 大きめに待っておく
    
    # 3) Python からリクエストを送る
    url = "http://localhost:8000/v1/chat/completions"
    payload = {
        "model": "deepseek-ai/DeepSeek-R1",
        "messages": [
            {"role": "user", "content": "What is the capital of France?"}
        ]
    }
    response = requests.post(url, json=payload)
    print(response.json())

    # 4) 用が済んだらサーバーを終了
    p.kill()

if __name__ == "__main__":
    main()

あるいは下記のように記述してもいいようだ。

main.py

from vllm import LLM, SamplingParams

def main():
    llm = LLM(model="deepseek-ai/DeepSeek-R1")
    prompt = "What is the capital of France?"
    outputs = llm.generate([prompt], SamplingParams(max_tokens=64))
    print(outputs[0].outputs[0].text)

if __name__ == "__main__":
    main()

vLLMを使わない方法

vLLMを使わない場合、下記のように書ける。

main.py

from transformers import pipeline

messages = [
    {"role": "user", "content": "Who are you?"},
]

# パイプラインの生成
pipe = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1", trust_remote_code=True)

# 推論結果を変数に代入
result = pipe(messages)

# 結果を表示
print(result)

また、pipelineを使わない場合、下記のように書けるらしい。

main.py

from transformers import AutoTokenizer, AutoModelForCausalLM

# トークナイザーの読み込み
tokenizer = AutoTokenizer.from_pretrained("deepseek-ai/DeepSeek-R1", trust_remote_code=True)
# モデルの読み込み
model = AutoModelForCausalLM.from_pretrained("deepseek-ai/DeepSeek-R1", trust_remote_code=True)

# 推論したい文字列
prompt = "Who are you?"

# トークナイズしてテンソル化
input_ids = tokenizer(prompt, return_tensors="pt").input_ids

# モデルを使ってトークン列を生成
output_ids = model.generate(input_ids, max_new_tokens=50)

# 生成されたトークン列を文字列にデコード
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)

# 結果の表示
print(output_text)

まとめ

明日会社で環境構築できるように頑張る