Gemma4 31Bのpure PyTorch実装（tokenizerのみtransformers使用）

gemma4_31b_minimal.py

"""
Gemma 4 31B Dense - Minimal Pure PyTorch Implementation (Text-only, Inference-only)

Architecture reference: huggingface/transformers modeling_gemma4.py
Target model: google/gemma-4-31B-it (or google/gemma-4-31B)

Design goals:
  - PyTorch only, no transformers dependency
  - Gemma 4 31B 固定 (config値をハードコード可)
  - 推論時の実験 (レイヤー繰り返し、logit調査等) がしやすい構造
"""

from pathlib import Path
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

# ============================================================================
# Constants from config.json (google/gemma-4-31B-it)
# ============================================================================

VOCAB_SIZE = 262144
HIDDEN_SIZE = 5376
NUM_LAYERS = 60
INTERMEDIATE_SIZE = 21504
RMS_NORM_EPS = 1e-6
LOGIT_SOFTCAP = 30.0

# Attention config - sliding layers
NUM_Q_HEADS = 32
NUM_KV_HEADS_SLIDING = 16
HEAD_DIM_SLIDING = 256
SLIDING_WINDOW = 1024
ROPE_THETA_SLIDING = 10000.0

# Attention config - full attention layers
NUM_KV_HEADS_FULL = 4  # num_global_key_value_heads
HEAD_DIM_FULL = 512  # global_head_dim
ROPE_THETA_FULL = 1000000.0
PARTIAL_ROTARY_FACTOR = 0.25  # full attentionでは head_dim の 25% (=128次元) のみ回転

# Layer pattern: 5 sliding + 1 full, repeated 10 times = 60 layers
LAYER_TYPES = (["sliding"] * 5 + ["full"]) * 10


# ============================================================================
# Basic components
# ============================================================================


class RMSNorm(nn.Module):
    """
    Gemma 4 の RMSNorm: weight * (x / rms(x))
    NOTE: Gemma 2 では (1 + weight) * norm(x) だったが、Gemma 4 では weight を直接乗算。
          weight は ones で初期化されるので、初期値での動作は同じ。
    """

    def __init__(self, dim: int, eps: float = RMS_NORM_EPS):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_float = x.float()
        normed = x_float * torch.pow(x_float.pow(2).mean(-1, keepdim=True) + self.eps, -0.5)
        return (normed * self.weight.float()).to(x.dtype)


class RMSNormNoScale(nn.Module):
    """
    RMSNorm without learnable scale (with_scale=False).
    Used for v_norm in attention.
    """

    def __init__(self, dim: int, eps: float = RMS_NORM_EPS):
        super().__init__()
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_float = x.float()
        normed = x_float * torch.pow(x_float.pow(2).mean(-1, keepdim=True) + self.eps, -0.5)
        return normed.to(x.dtype)


class ScaledEmbedding(nn.Embedding):
    """
    Gemma 系の embedding: 出力を sqrt(hidden_size) でスケール。
    bfloat16 にダウンキャストしてからスケールする (transformers のコメント参照)。
    """

    def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int = 0):
        super().__init__(num_embeddings, embedding_dim, padding_idx=padding_idx)
        # bfloat16 でのスケール値 (sqrt(5376) ≈ 73.32... → bf16 で丸められる)
        self.register_buffer("embed_scale", torch.tensor(float(embedding_dim**0.5)), persistent=False)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        return super().forward(input_ids) * self.embed_scale.to(self.weight.dtype)


# ============================================================================
# RoPE (Rotary Position Embedding)
# ============================================================================


def build_rope_cache(
    seq_len: int,
    head_dim: int,
    theta: float,
    device: torch.device,
    partial_rotary_factor: float = 1.0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    RoPE の cos/sin テーブルを事前計算。

    partial_rotary_factor < 1.0 の場合、proportional RoPE を使用:
    - inv_freq の周波数計算は head_dim 全体で割る（rotary_dim ではなく）
    - 回転しない次元には inv_freq=0 を入れ、cos=1, sin=0 として恒等にする
    - cos/sin は head_dim 全体のサイズで返す

    Returns:
        cos, sin: shape (seq_len, head_dim)
    """
    rope_angles = int(partial_rotary_factor * head_dim // 2)
    # inv_freq の計算は常に head_dim で割る（proportional RoPE の仕様）
    inv_freq_rotated = 1.0 / (theta ** (torch.arange(0, 2 * rope_angles, 2, device=device, dtype=torch.float32) / head_dim))
    # 回転しない次元にはゼロを入れる（cos=1, sin=0 → 恒等変換）
    nope_angles = head_dim // 2 - rope_angles
    if nope_angles > 0:
        inv_freq = torch.cat([inv_freq_rotated, torch.zeros(nope_angles, device=device, dtype=torch.float32)])
    else:
        inv_freq = inv_freq_rotated

    # positions: (seq_len,)
    positions = torch.arange(seq_len, device=device, dtype=torch.float32)
    # freqs: (seq_len, head_dim // 2)
    freqs = torch.outer(positions, inv_freq)
    # emb: (seq_len, head_dim)
    emb = torch.cat([freqs, freqs], dim=-1)
    return emb.cos(), emb.sin()


def apply_rope(
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    position_ids: torch.Tensor,
) -> torch.Tensor:
    """
    x: (batch, seq_len, num_heads, head_dim)
    cos, sin: (max_seq_len, head_dim)  ← build_rope_cache の出力
    position_ids: (batch, seq_len)

    cos/sin は常に head_dim 全体のサイズ。proportional RoPE の場合、
    回転しない次元は cos=1, sin=0 なので自動的に恒等変換になる。

    Returns: same shape as x
    """
    # Gather cos/sin for requested positions: (batch, seq_len, head_dim)
    cos_pos = cos[position_ids]
    sin_pos = sin[position_ids]
    # Unsqueeze for head dimension: (batch, seq_len, 1, head_dim)
    cos_pos = cos_pos.unsqueeze(2)
    sin_pos = sin_pos.unsqueeze(2)

    # rotate_half: (x1, x2) → (-x2, x1)
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    x_rotated = torch.cat([-x2, x1], dim=-1)

    out = x * cos_pos + x_rotated * sin_pos
    return out.to(x.dtype)


# ============================================================================
# Attention
# ============================================================================


class Gemma4Attention(nn.Module):
    """
    Sliding / Full attention を統合した Attention モジュール。

    Key differences between sliding and full:
      - Sliding: head_dim=256, 16 KV heads, standard RoPE (theta=10k), causal + sliding window
      - Full:    head_dim=512, 4 KV heads, proportional RoPE (theta=1M, partial=0.25),
                 causal (no window limit), K=V (key states are reused as values)
    """

    def __init__(self, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.is_sliding = LAYER_TYPES[layer_idx] == "sliding"

        # Sliding vs Full で異なるパラメータ
        if self.is_sliding:
            self.head_dim = HEAD_DIM_SLIDING
            self.num_kv_heads = NUM_KV_HEADS_SLIDING
            self.k_eq_v = False
        else:
            self.head_dim = HEAD_DIM_FULL
            self.num_kv_heads = NUM_KV_HEADS_FULL
            self.k_eq_v = True  # attention_k_eq_v: full attention layers only

        self.num_q_heads = NUM_Q_HEADS
        self.num_kv_groups = self.num_q_heads // self.num_kv_heads

        # Projections
        self.q_proj = nn.Linear(HIDDEN_SIZE, self.num_q_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(HIDDEN_SIZE, self.num_kv_heads * self.head_dim, bias=False)
        # v_proj: k_eq_v=True のレイヤーでは None（key states を value として再利用）
        if not self.k_eq_v:
            self.v_proj = nn.Linear(HIDDEN_SIZE, self.num_kv_heads * self.head_dim, bias=False)
        else:
            self.v_proj = None
        self.o_proj = nn.Linear(self.num_q_heads * self.head_dim, HIDDEN_SIZE, bias=False)

        # QK norms (per-head normalization, with learnable scale)
        self.q_norm = RMSNorm(self.head_dim)
        self.k_norm = RMSNorm(self.head_dim)
        # V norm (without learnable scale - just normalization)
        self.v_norm = RMSNormNoScale(self.head_dim)

    def forward(
        self,
        hidden_states: torch.Tensor,  # (batch, seq_len, hidden_size)
        position_ids: torch.Tensor,  # (batch, seq_len)
        cos: torch.Tensor,  # RoPE cos cache
        sin: torch.Tensor,  # RoPE sin cache
        kv_cache: Optional[dict] = None,  # {"key": Tensor, "value": Tensor} in (B, heads, kv_len, dim)
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        bsz, q_len, _ = hidden_states.shape

        # --- Q, K, V projection ---
        q = self.q_proj(hidden_states).view(bsz, q_len, self.num_q_heads, self.head_dim)
        k = self.k_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim)

        # k_eq_v: value は k_proj の出力を直接使う（norm/RoPE 前のテンソル）
        if self.v_proj is not None:
            v = self.v_proj(hidden_states).view(bsz, q_len, self.num_kv_heads, self.head_dim)
        else:
            v = k  # k_eq_v: pre-norm, pre-RoPE の k_proj 出力

        # QK norms (applied per head, before RoPE)
        q = self.q_norm(q)
        k = self.k_norm(k)

        # RoPE (Q と K のみ。V には適用しない)
        q = apply_rope(q, cos, sin, position_ids)
        k = apply_rope(k, cos, sin, position_ids)

        # V norm (learnable scale なし、全レイヤー共通)
        v = self.v_norm(v)

        # --- Transpose to (batch, heads, seq, dim) for attention ---
        q = q.transpose(1, 2)  # (bsz, num_q_heads, q_len, head_dim)
        k = k.transpose(1, 2)  # (bsz, num_kv_heads, q_len, head_dim)
        v = v.transpose(1, 2)

        # --- KV cache: concat with past, then store ---
        if kv_cache is not None:
            past_k = kv_cache.get("key")
            if past_k is not None:
                k = torch.cat([past_k, k], dim=2)  # (bsz, heads, past+q_len, dim)
                v = torch.cat([kv_cache["value"], v], dim=2)

            # Sliding window レイヤーでは最新 SLIDING_WINDOW エントリのみ保持
            if self.is_sliding and k.shape[2] > SLIDING_WINDOW:
                k = k[:, :, -SLIDING_WINDOW:]
                v = v[:, :, -SLIDING_WINDOW:]

            kv_cache["key"] = k
            kv_cache["value"] = v

        kv_len = k.shape[2]

        # --- GQA: expand KV heads ---
        if self.num_kv_groups > 1:
            k = k.repeat_interleave(self.num_kv_groups, dim=1)
            v = v.repeat_interleave(self.num_kv_groups, dim=1)

        # --- Attention ---
        # NOTE: Gemma 4 uses scaling=1.0 (QK norms control magnitude)
        # NOTE: No attention logit softcapping (final_logit_softcapping is for LM head only)
        #
        # Masking strategy:
        #   - q_len > 1 (prefill): causal mask needed
        #   - q_len == 1 (decode): single query attends to all KV, no mask needed
        if attention_mask is not None:
            attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, scale=1.0)
        elif q_len > 1:
            attn_output = F.scaled_dot_product_attention(q, k, v, is_causal=True, scale=1.0)
        else:
            attn_output = F.scaled_dot_product_attention(q, k, v, scale=1.0)

        # --- Output projection ---
        attn_output = attn_output.transpose(1, 2).reshape(bsz, q_len, -1)
        return self.o_proj(attn_output)


# ============================================================================
# MLP (GeGLU)
# ============================================================================


class Gemma4MLP(nn.Module):
    """
    GeGLU: gate_proj で GeLU を通し、up_proj と要素積を取り、down_proj で射影。
    hidden_activation = "gelu_pytorch_tanh" → F.gelu(x, approximate='tanh')
    """

    def __init__(self):
        super().__init__()
        self.gate_proj = nn.Linear(HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False)
        self.up_proj = nn.Linear(HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False)
        self.down_proj = nn.Linear(INTERMEDIATE_SIZE, HIDDEN_SIZE, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.gelu(self.gate_proj(x), approximate="tanh") * self.up_proj(x))


# ============================================================================
# Decoder Layer
# ============================================================================


class Gemma4DecoderLayer(nn.Module):
    """
    Sandwich norm パターン: attention / MLP の前後に RMSNorm。
    レイヤー末尾で layer_scalar を乗算。

    Forward:
        residual = x
        x = input_layernorm(x)
        x = attention(x)
        x = post_attention_layernorm(x)
        x = residual + x

        residual = x
        x = pre_feedforward_layernorm(x)
        x = mlp(x)
        x = post_feedforward_layernorm(x)
        x = residual + x

        x = x * layer_scalar
    """

    def __init__(self, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.self_attn = Gemma4Attention(layer_idx)
        self.mlp = Gemma4MLP()

        # Sandwich norms (4 per layer)
        self.input_layernorm = RMSNorm(HIDDEN_SIZE)
        self.post_attention_layernorm = RMSNorm(HIDDEN_SIZE)
        self.pre_feedforward_layernorm = RMSNorm(HIDDEN_SIZE)
        self.post_feedforward_layernorm = RMSNorm(HIDDEN_SIZE)

        # Per-layer learned scalar (buffer, not parameter — 学習時は凍結?)
        self.register_buffer("layer_scalar", torch.ones(1))

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_ids: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
        kv_cache: Optional[dict] = None,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        # --- Attention block ---
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states = self.self_attn(hidden_states, position_ids, cos, sin, kv_cache, attention_mask)
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = residual + hidden_states

        # --- MLP block ---
        residual = hidden_states
        hidden_states = self.pre_feedforward_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = self.post_feedforward_layernorm(hidden_states)
        hidden_states = residual + hidden_states

        # --- Per-layer scaling ---
        hidden_states = hidden_states * self.layer_scalar

        return hidden_states


# ============================================================================
# Full Model
# ============================================================================


class Gemma4TextModel(nn.Module):
    """
    Gemma 4 31B Dense text decoder.

    tie_word_embeddings=True なので lm_head は embed_tokens.weight を共有。
    最終 logits には softcapping (tanh * 30.0) を適用。
    """

    def __init__(self):
        super().__init__()
        self.embed_tokens = ScaledEmbedding(VOCAB_SIZE, HIDDEN_SIZE)
        self.layers = nn.ModuleList([Gemma4DecoderLayer(i) for i in range(NUM_LAYERS)])
        self.norm = RMSNorm(HIDDEN_SIZE)  # final norm

        # RoPE caches (will be initialized on first forward)
        self._rope_sliding: Optional[tuple[torch.Tensor, torch.Tensor]] = None
        self._rope_full: Optional[tuple[torch.Tensor, torch.Tensor]] = None

    def _ensure_rope_cache(self, max_seq_len: int, device: torch.device):
        """RoPE cos/sin テーブルを必要に応じて生成・拡張。"""
        if self._rope_sliding is None or self._rope_sliding[0].shape[0] < max_seq_len:
            self._rope_sliding = build_rope_cache(max_seq_len, HEAD_DIM_SLIDING, ROPE_THETA_SLIDING, device)
        if self._rope_full is None or self._rope_full[0].shape[0] < max_seq_len:
            self._rope_full = build_rope_cache(
                max_seq_len,
                HEAD_DIM_FULL,
                ROPE_THETA_FULL,
                device,
                partial_rotary_factor=PARTIAL_ROTARY_FACTOR,
            )

    def forward(
        self,
        input_ids: torch.Tensor,  # (batch, seq_len)
        position_ids: Optional[torch.Tensor] = None,  # (batch, seq_len)
        kv_caches: Optional[list[dict]] = None,  # per-layer KV cache
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Returns: logits (batch, seq_len, vocab_size) — softcap 適用済み

        KV cache 使用時:
          - kv_caches は長さ NUM_LAYERS のリスト。各要素は {} (初回) または {"key": ..., "value": ...}
          - position_ids を指定しない場合、cache 内の past_len から自動計算
        """
        bsz, seq_len = input_ids.shape
        device = input_ids.device

        if position_ids is None:
            # KV cache から past_len を推定 (full attention レイヤーの cache を参照)
            past_len = 0
            if kv_caches is not None:
                # full attention レイヤーは eviction しないので正確な past_len が取れる
                for i, c in enumerate(kv_caches):
                    if not LAYER_TYPES[i] == "sliding" and "key" in c:
                        past_len = c["key"].shape[2]
                        break
                else:
                    # full attention に cache がなければ sliding を確認
                    for c in kv_caches:
                        if "key" in c:
                            past_len = c["key"].shape[2]
                            break
            position_ids = torch.arange(past_len, past_len + seq_len, device=device).unsqueeze(0).expand(bsz, -1)

        # RoPE cache
        max_pos = position_ids.max().item() + 1
        self._ensure_rope_cache(max_pos, device)
        cos_sliding, sin_sliding = self._rope_sliding
        cos_full, sin_full = self._rope_full

        # Embedding
        hidden_states = self.embed_tokens(input_ids)

        # Decoder layers
        for i, layer in enumerate(self.layers):
            is_sliding = LAYER_TYPES[i] == "sliding"
            cos = cos_sliding if is_sliding else cos_full
            sin = sin_sliding if is_sliding else sin_full
            kv_cache = kv_caches[i] if kv_caches is not None else None

            # TODO: attention_mask をレイヤー種別に応じて切り替え
            #   - sliding: causal + sliding window (最新1024トークンのみ)
            #   - full: causal (全トークン)

            hidden_states = layer(hidden_states, position_ids, cos, sin, kv_cache, attention_mask)

        # Final norm
        hidden_states = self.norm(hidden_states)

        # LM head (tied weights)
        logits = F.linear(hidden_states, self.embed_tokens.weight)

        # Logit softcapping: tanh(logits / cap) * cap
        logits = torch.tanh(logits / LOGIT_SOFTCAP) * LOGIT_SOFTCAP

        return logits


# ============================================================================
# Weight Loading
# ============================================================================


def load_weights(model: Gemma4TextModel, model_dir: str | Path):
    """
    safetensors から重みを読み込む。

    Weight name mapping (safetensors → this model):
      model.language_model.embed_tokens.weight    → embed_tokens.weight
      model.language_model.layers.{i}.xxx         → layers.{i}.xxx
      model.language_model.norm.weight             → norm.weight

    NOTE: vision_tower / embed_vision は無視 (text-only)
    NOTE: k_eq_v レイヤーの v_proj 重みはファイルに存在するが読み込まない
    """
    from safetensors import safe_open

    model_dir = Path(model_dir)

    # Find all safetensor files
    safetensor_files = sorted(model_dir.glob("model*.safetensors"))
    if not safetensor_files:
        raise FileNotFoundError(f"No safetensor files found in {model_dir}")

    # Build mapping: our param name → safetensors key name
    prefix = "model.language_model."
    state_dict = model.state_dict()

    loaded = set()
    for sf_path in safetensor_files:
        with safe_open(sf_path, framework="pt") as f:
            for sf_key in f.keys():
                # Skip non-language-model weights
                if not sf_key.startswith(prefix):
                    continue

                our_key = sf_key[len(prefix) :]

                # Skip v_proj for k_eq_v layers (full attention)
                # layer index を抽出して判定
                if "self_attn.v_proj" in our_key:
                    parts = our_key.split(".")
                    layer_idx = int(parts[1])
                    if LAYER_TYPES[layer_idx] == "full":
                        continue  # k_eq_v: v_proj は不要

                if our_key in state_dict:
                    tensor = f.get_tensor(sf_key)
                    if state_dict[our_key].shape != tensor.shape:
                        print(f"  Shape mismatch: {our_key} " f"expected {state_dict[our_key].shape}, " f"got {tensor.shape}")
                        continue
                    state_dict[our_key] = tensor
                    loaded.add(our_key)

    # Report
    missing = set(state_dict.keys()) - loaded
    if missing:
        print(f"Missing weights ({len(missing)}):")
        for k in sorted(missing):
            print(f"  {k}")

    # assign=Trueでmeta tensorを直接置き換える。ただし、safetensorsはlazy loading
    model.load_state_dict(state_dict, strict=False, assign=True)
    print(f"Loaded {len(loaded)} / {len(state_dict)} parameters")


# ============================================================================
# Simple text generation (greedy)
# ============================================================================


@torch.no_grad()
def generate(
    model: Gemma4TextModel,
    input_ids: torch.Tensor,
    max_new_tokens: int = 100,
    eos_token_id: int | list[int] = 1,
    temperature: float = 1.0,
) -> torch.Tensor:
    """
    Greedy / temperature sampling による生成 (KV cache 使用)。

    1. Prefill: 入力全体を一度に処理し、KV cache を構築
    2. Decode: 1トークンずつ生成、cache を更新
    """
    if isinstance(eos_token_id, int):
        eos_token_id = [eos_token_id]

    # KV cache を初期化（空の dict のリスト）
    kv_caches: list[dict] = [{} for _ in range(NUM_LAYERS)]

    # --- Prefill: 入力プロンプト全体を処理 ---
    logits = model(input_ids, kv_caches=kv_caches)
    next_logits = logits[:, -1, :]

    generated = input_ids.clone()

    for _ in range(max_new_tokens):
        if temperature > 0 and temperature != 1.0:
            next_logits = next_logits / temperature

        # --- ここで logits を自由に調査・加工可能 ---
        # 例: top-k, top-p, repetition penalty, logit lens, etc.

        next_token = next_logits.argmax(dim=-1, keepdim=True)  # greedy
        generated = torch.cat([generated, next_token], dim=1)

        # EOS check
        if next_token.item() in eos_token_id:
            break

        # --- Decode: 新しい1トークンだけ処理 ---
        logits = model(next_token, kv_caches=kv_caches)
        next_logits = logits[:, -1, :]

    return generated


# ============================================================================
# Entry point (example usage)
# ============================================================================

if __name__ == "__main__":
    import sys

    # Hugging Face Hub からダウンロードしたモデルの safetensors ディレクトリを指定
    model_dir = sys.argv[1] if len(sys.argv) > 1 else "./gemma-4-31B-it"
    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.bfloat16

    print("Building model...")

    # TODO init_empty_weightsのためだけにaccelerateを入れるのは微妙なので、将来的には自前で同様の機能を実装しても良いかも
    from accelerate import init_empty_weights

    # init_empty_weightsを使わないと31*4=124GB以上のVRAMが必要になるため、モデル構造だけを先に定義してから重みをロードする
    with init_empty_weights():
        model = Gemma4TextModel()
    # model.to(dtype=dtype, device=device)

    print(f"Loading weights from {model_dir}...")
    load_weights(model, model_dir)

    model.to(dtype=dtype, device=device)  # ここで実際のロードが起こる
    model.eval()

    # # Simple test
    # # tokenizer は sentencepiece を直接使うか、transformers.AutoTokenizer を使う
    # # ここでは input_ids を直接指定する例
    # test_ids = torch.tensor([[2, 1645, 108]], device=device)  # BOS + example
    # print("Generating...")
    # with torch.no_grad(), torch.autocast(device_type=device, dtype=dtype, enabled=True):
    #     output = generate(model, test_ids, max_new_tokens=50)
    # print(f"Output token ids: {output[0].tolist()}")

    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(model_dir)

    msgs = [{"role": "user", "content": "こんにちは、調子はどう？"}]
    prompt = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True, enable_thinking=True)  # False)
    input_ids = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
    print(f"Input tokens ({input_ids.shape[1]}): {input_ids[0].tolist()}")
    print(f"Decoded: {tokenizer.decode(input_ids[0])}")

    print("\nGenerating...")
    import time

    start_time = time.perf_counter()
    with torch.no_grad():
        output = generate(model, input_ids, max_new_tokens=1000, eos_token_id=[1, 106])
    end_time = time.perf_counter()
    print(f"\nGeneration completed in {end_time - start_time:.2f} seconds.")

    output_ids = output[0].tolist()
    generated_ids = output_ids[input_ids.shape[1] :]
    print(f"\nGeneration speed: {len(generated_ids) / (end_time - start_time):.2f} tokens/sec")
    print(f"\nGenerated token ids: {generated_ids[:30]}...")
    print(f"\nDecoded output:\n{tokenizer.decode(generated_ids, skip_special_tokens=False)}")