圧縮バージョンのDeepSeekをローカルで。CPUとDRAMで推論,GPUとVRAMで推論。超簡易備忘録。

import

from transformers import AutoTokenizer
from transformers import AutoModelForCausalLM
import torch

完成関数（CPUとDRAMで推論バージョン）

## CPU-DRAMでの推論は、入力コーパス500の出力300では、1時間40分程度かかります。

        def run_deepseek_inference(self, corpus: str) -> str:
            model_name = "deepseek-ai/deepseek-coder-1.3b-base"

            tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
            tokenizer.padding_side = "left" 

            model = AutoModelForCausalLM.from_pretrained(
                model_name,
                device_map=None,  # GPU使わない
                torch_dtype=torch.float32,
                low_cpu_mem_usage=True,
                use_auth_token=True,
            )

            inputs = tokenizer(corpus, return_tensors="pt")
            max_new_tokens = 300  # ここで最大トークン数を指定
            generated_text = corpus  # 初期入力はコーパスそのまま

            # 進捗バーの初期化
            progress_bar = tqdm(range(max_new_tokens), desc="Generating", unit="token")

            for _ in progress_bar:
                # モデルで1トークンずつ生成
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=1,
                    eos_token_id=None,
                    pad_token_id=tokenizer.eos_token_id
                )
                
                # 新しく生成されたトークンをデコードして追加
                new_token = tokenizer.decode(outputs[0, -1:], skip_special_tokens=True)
                generated_text += new_token
                
                # 次回の入力には生成したテキストを追加
                inputs = tokenizer(generated_text, return_tensors="pt")

                # 進捗バーを更新
                progress_bar.set_postfix({"Output": new_token})

            return generated_text

完成関数_2（GPUとVRAMで推論バージョン）（3060でも全然イケちゃった）

## GPU-VRAMでの推論は、入力コーパス500の出力300では、4分程度かかります。
        def run_deepseek_inference(self, corpus: str) -> str:
            model_name = "deepseek-ai/deepseek-coder-1.3b-base"

            tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
            tokenizer.padding_side = "left" 

            model = AutoModelForCausalLM.from_pretrained(
                model_name,
                device_map="auto",  # ← GPU使用に変更！
                torch_dtype=torch.float16,  # GPUでの高速処理に適した精度
                low_cpu_mem_usage=True,
                use_auth_token=True,
            )
            model.eval()  # 推論モードに

            generated_text = corpus
            max_new_tokens = 300

            # 初回トークン化＋GPUへ移動
            inputs = tokenizer(generated_text, return_tensors="pt").to("cuda")

            progress_bar = tqdm(range(max_new_tokens), desc="Generating", unit="token")

            for _ in progress_bar:
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=1,
                    eos_token_id=None,
                    pad_token_id=tokenizer.eos_token_id
                )

                # GPU → CPU に戻してデコード
                new_token = tokenizer.decode(outputs[0, -1:].cpu(), skip_special_tokens=True)
                generated_text += new_token

                # 再トークン化＋GPUに再配置
                inputs = tokenizer(generated_text, return_tensors="pt").to("cuda")

                progress_bar.set_postfix({"Output": new_token})

            return generated_text

入力サイズと出力サイズのおすすめ。（chatGPT4oより。）

余談

尚、流石にpython3.6のtransformersライブラリ内にはDeepSeekはサポートされておらず、python3.9での実装となった。