inference.py

# ==============================================================================
#                          Inference Script
# ==============================================================================
# --- Necessary Imports ---
import torch
import torch.nn as nn
from dataclasses import dataclass, field
import math
import torch.nn.functional as F
from transformers import AutoTokenizer
import os
import glob
import time
import datetime
import traceback
import dataclasses # Make sure this is imported

# --- Model Configuration ---
# IMPORTANT: This definition MUST exactly match the one used during training
#            when the checkpoint was saved.
@dataclass
class ModelArgs:
    # --- ~221M Config used for training step_1200 ---
    hidden_size: int = 768; num_hidden_layers: int = 12; num_attention_heads: int = 12
    num_key_value_heads: int = 12; intermediate_size: int = 2048; vocab_size: int = 128000
    rms_norm_eps: float = 1e-5; rope_theta: float = 500000.0; max_position_embeddings: int = 4096
    head_dim: int = field(init=False)
    add_recency_bias: bool = False # Ensure this matches the value used when saving the checkpoint

    def __post_init__(self):
        self.head_dim = self.hidden_size // self.num_attention_heads
        if self.hidden_size % self.num_attention_heads != 0: raise ValueError("hidden_size % num_attention_heads != 0")
        if self.num_attention_heads % self.num_key_value_heads != 0: raise ValueError("num_attention_heads % num_key_value_heads != 0")

# --- Model Components (RMSNorm, RoPE funcs, Attention, FeedForward, TransformerBlock, Llama) ---
# V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V V
# --- PASTE THE FULL DEFINITIONS OF THE FOLLOWING CLASSES/FUNCTIONS HERE ---
# --- from your model_architecture.py script:                         ---
#
class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6): super().__init__(); self.eps = eps; self.weight = nn.Parameter(torch.ones(dim))
    def _norm(self, x): return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
    def forward(self, x): original_dtype = x.dtype; output = self._norm(x.float()).to(original_dtype); return output * self.weight

def precompute_theta_pos_frequencies(head_dim: int, seq_len: int, device: str | torch.device, theta: float = 10000.0):
    if head_dim % 2 != 0: raise ValueError("head_dim must be even for RoPE")
    theta_indices = torch.arange(0, head_dim, 2).float(); theta_freqs = 1.0 / (theta**(theta_indices / head_dim))
    target_device = torch.device(device) if isinstance(device, str) else device; theta_freqs = theta_freqs.to(target_device)
    positions = torch.arange(seq_len, device=target_device).float(); freqs = torch.outer(positions, theta_freqs).float(); return freqs, positions

def apply_rotary_embeddings(x: torch.Tensor, freqs_cis_full: torch.Tensor, positions: torch.Tensor):
    positions = positions.long(); max_pos = freqs_cis_full.shape[0]
    if torch.max(positions) >= max_pos: positions = torch.clamp(positions, max=max_pos - 1)
    freqs = freqs_cis_full[positions]; freqs = freqs.unsqueeze(0).unsqueeze(2)
    bsz, seq_len, n_part_heads, head_dim = x.shape; x1 = x[..., : head_dim // 2]; x2 = x[..., head_dim // 2 :]
    cos_freqs = torch.cos(freqs).type_as(x); sin_freqs = torch.sin(freqs).type_as(x)
    rotated_x1 = x1 * cos_freqs - x2 * sin_freqs; rotated_x2 = x1 * sin_freqs + x2 * cos_freqs
    rotated_x = torch.cat([rotated_x1, rotated_x2], dim=-1); return rotated_x.type_as(x)

class Attention(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__(); self.args = args; self.num_heads = args.num_attention_heads; self.num_kv_heads = args.num_key_value_heads
        self.head_dim = args.head_dim; self.repeats = self.num_heads // self.num_kv_heads
        self.wq = nn.Linear(args.hidden_size, args.num_attention_heads * args.head_dim, bias=False)
        self.wk = nn.Linear(args.hidden_size, args.num_key_value_heads * args.head_dim, bias=False)
        self.wv = nn.Linear(args.hidden_size, args.num_key_value_heads * args.head_dim, bias=False)
        self.wo = nn.Linear(args.num_attention_heads * args.head_dim, args.hidden_size, bias=False)
    def _repeat_kv(self, x: torch.Tensor, n_rep: int) -> torch.Tensor:
        bsz, n_kv_heads, seqlen, head_dim = x.shape;
        if n_rep == 1: return x
        return (x[:, :, None, :, :].expand(bsz, n_kv_heads, n_rep, seqlen, head_dim).reshape(bsz, n_kv_heads * n_rep, seqlen, head_dim))
    def _create_recency_bias(self, seqlen, full_seqlen, device, dtype, bias_strength=0.1, decay_rate=0.9):
        bias = torch.zeros((1, 1, seqlen, full_seqlen), device=device, dtype=dtype); indices = torch.arange(full_seqlen, device=device)
        rel_pos = torch.arange(seqlen, device=device).unsqueeze(1) - indices.unsqueeze(0); mask = rel_pos >= 0
        decaying_bias = bias_strength * (decay_rate ** (-rel_pos[mask])); bias[:, :, mask] = decaying_bias.type_as(bias); return bias
    def forward(self, x: torch.Tensor, freqs_cis_full: torch.Tensor, positions: torch.Tensor, mask: torch.Tensor | None = None, cache: tuple[torch.Tensor, torch.Tensor] | None = None) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        bsz, seqlen, _ = x.shape; xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
        xq = xq.view(bsz, seqlen, self.num_heads, self.head_dim); xk = xk.view(bsz, seqlen, self.num_kv_heads, self.head_dim); xv = xv.view(bsz, seqlen, self.num_kv_heads, self.head_dim)
        xq = apply_rotary_embeddings(xq, freqs_cis_full, positions); xk = apply_rotary_embeddings(xk, freqs_cis_full, positions)
        xk = xk.transpose(1, 2); xv = xv.transpose(1, 2)
        if cache is not None: cache_k, cache_v = cache; keys = torch.cat((cache_k.to(xk.device), xk), dim=2); values = torch.cat((cache_v.to(xv.device), xv), dim=2)
        else: keys = xk; values = xv
        updated_cache = (keys.detach(), values.detach()); keys_repeated = self._repeat_kv(keys, self.repeats); values_repeated = self._repeat_kv(values, self.repeats)
        xq = xq.transpose(1, 2); scores = torch.matmul(xq.float(), keys_repeated.transpose(-2, -1).float()) / math.sqrt(self.head_dim)
        if self.args.add_recency_bias:
             full_seqlen = keys_repeated.shape[-2]; recency_bias = self._create_recency_bias(seqlen, full_seqlen, device=scores.device, dtype=scores.dtype); scores = scores + recency_bias
        if mask is not None:
            full_seqlen = keys_repeated.shape[-2]; expected_mask_shape_end = (seqlen, full_seqlen)
            if mask.shape[-2:] != expected_mask_shape_end:
                 try: mask_slice = mask[:, :, -seqlen:, :full_seqlen]; scores = scores + mask_slice.float()
                 except Exception: pass
            else: scores = scores + mask.float()
        scores = nn.functional.softmax(scores, dim=-1).type_as(xq); output = torch.matmul(scores, values_repeated)
        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1); output = self.wo(output); return output, updated_cache

class FeedForward(nn.Module):
    def __init__(self, args: ModelArgs): super().__init__(); self.w1 = nn.Linear(args.hidden_size, args.intermediate_size, bias=False); self.w3 = nn.Linear(args.hidden_size, args.intermediate_size, bias=False); self.w2 = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
    def forward(self, x: torch.Tensor) -> torch.Tensor: return self.w2(nn.functional.silu(self.w1(x)) * self.w3(x))

class TransformerBlock(nn.Module):
    def __init__(self, args: ModelArgs): super().__init__(); self.args = args; self.attention_norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps); self.attention = Attention(args); self.ffn_norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps); self.feed_forward = FeedForward(args)
    def forward(self, x: torch.Tensor, freqs_cis_full: torch.Tensor, positions: torch.Tensor, mask: torch.Tensor | None = None, cache: tuple[torch.Tensor, torch.Tensor] | None = None) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
        r, cache = self.attention(self.attention_norm(x), freqs_cis_full, positions, mask, cache); h = x + r; r = self.feed_forward(self.ffn_norm(h)); out = h + r; return out, cache

class Llama(nn.Module):
    def __init__(self, args: ModelArgs):
        super().__init__(); self.args = args; self.tok_embeddings = nn.Embedding(args.vocab_size, args.hidden_size); self.layers = nn.ModuleList([TransformerBlock(args) for _ in range(args.num_hidden_layers)])
        self.norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps); self.tok_embeddings.weight.requires_grad = True
        freqs_cis, _ = precompute_theta_pos_frequencies(args.head_dim, args.max_position_embeddings, device='cpu', theta=args.rope_theta)
        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
    def forward(self, tokens: torch.Tensor, positions: torch.Tensor):
        bsz, seqlen = tokens.shape; h = self.tok_embeddings(tokens); freqs_cis_full = self.freqs_cis.to(h.device); mask = None
        if seqlen > 1: mask = torch.full((1, 1, seqlen, seqlen), float("-inf"), device=h.device); mask = torch.triu(mask, diagonal=1).type_as(h)
        positions = positions.to(h.device)
        for layer in self.layers: h, _ = layer(h, freqs_cis_full, positions, mask, cache=None) # Pass cache=None for non-cached forward
        h = self.norm(h); output = F.linear(h, self.tok_embeddings.weight); return output # Use tied weights

@torch.no_grad()
def generate(model: Llama, tokenizer: AutoTokenizer, prompt: str, max_new_tokens: int, temperature: float = 1.0, top_k: int | None = None, top_p: float | None = None):
    model.eval() # CORRECTED: Separate line
    try:
        model_device = next(model.parameters()).device
        model_dtype = next(model.parameters()).dtype
    except StopIteration:
        print("Warning: Model has no parameters. Assuming CPU and float32.")
        model_device = torch.device("cpu")
        model_dtype = torch.float32

    prompt_ids = tokenizer.encode(prompt, add_special_tokens=True); tokens = torch.tensor([prompt_ids], dtype=torch.long, device=model_device)
    cache = [(torch.zeros((1, model.args.num_key_value_heads, 0, model.args.head_dim), device=model_device, dtype=model_dtype),
              torch.zeros((1, model.args.num_key_value_heads, 0, model.args.head_dim), device=model_device, dtype=model_dtype))
             for _ in range(model.args.num_hidden_layers)]
    generated_token_ids = []; current_tokens = tokens; print(f"Generating {max_new_tokens} tokens from prompt: '{prompt}'"); print("Output: ", end='')
    full_freqs_cis = model.freqs_cis.to(model_device)
    for i in range(max_new_tokens):
        current_seq_len = current_tokens.shape[1]; start_pos = cache[0][0].shape[2]; positions = torch.arange(start_pos, start_pos + current_seq_len, device=model_device)
        current_mask = None;
        if i == 0 and current_seq_len > 1: current_mask = torch.full((1, 1, current_seq_len, current_seq_len), float("-inf"), device=model_device); current_mask = torch.triu(current_mask, diagonal=1).type(model_dtype)
        h = model.tok_embeddings(current_tokens); updated_cache_list = []
        for layer_idx, layer in enumerate(model.layers): h, updated_layer_cache = layer(h, full_freqs_cis, positions, current_mask, cache[layer_idx]); updated_cache_list.append(updated_layer_cache)
        cache = updated_cache_list; h = model.norm(h); logits = F.linear(h, model.tok_embeddings.weight)
        next_token_logits = logits[:, -1, :]
        if temperature == 0: next_token_id = torch.argmax(next_token_logits, dim=-1, keepdim=True)
        else:
            next_token_logits = next_token_logits / temperature
            if top_k is not None and top_k > 0: v, _ = torch.topk(next_token_logits, min(top_k, next_token_logits.size(-1))); next_token_logits[next_token_logits < v[:, [-1]]] = float('-inf')
            if top_p is not None and 0.0 < top_p < 1.0:
                probs_for_filter = F.softmax(next_token_logits, dim=-1); probs_sort, probs_idx = torch.sort(probs_for_filter, descending=True); probs_sum = torch.cumsum(probs_sort, dim=-1)
                mask_top_p = probs_sum > top_p; mask_top_p[..., 0] = False; mask_top_p[..., 1:] = mask_top_p[..., :-1].clone(); indices_to_remove = mask_top_p.scatter(1, probs_idx, mask_top_p); next_token_logits[indices_to_remove] = float('-inf')
            probs = F.softmax(next_token_logits, dim=-1); next_token_id = torch.multinomial(probs, num_samples=1)
        if tokenizer.eos_token_id is not None and next_token_id.item() == tokenizer.eos_token_id: print("\n[EOS token reached]"); break
        next_token_id_item = next_token_id.item(); generated_token_ids.append(next_token_id_item); current_tokens = next_token_id.clone()
        print(tokenizer.decode([next_token_id_item]), end='', flush=True)
        if len(generated_token_ids) >= max_new_tokens: break
    print("\n--- Generation Complete ---"); final_token_ids = prompt_ids + generated_token_ids; full_generated_text = tokenizer.decode(final_token_ids, skip_special_tokens=False)
    print(f"\nFull generated text:\n{full_generated_text}"); return full_generated_text
# --- End Placeholders ---


# --- Main Inference Execution ---
if __name__ == "__main__":

    # --- Configuration for Inference ---
    # --- !! USE SPECIFIC WINDOWS PATH !! ---
    raw_checkpoint_path = r".\step_800.pt" # <<< CHANGED to step 1200
    # --- Normalize the path ---
    checkpoint_path = os.path.normpath(raw_checkpoint_path)
    # --- End Adjust ---

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"\n--- Inference Setup ---")
    print(f"Using device: {device}")
    print(f"Attempting to load checkpoint: {checkpoint_path}")

    # --- Load Checkpoint and Model Args ---
    if not os.path.exists(checkpoint_path):
        # Removed the fallback logic as we are specifying an exact path
        exit(f"Error: Checkpoint file not found at the specified path: {checkpoint_path}")

    try:
        # Load checkpoint to CPU first
        checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False) # weights_only=False needed for ModelArgs

        # Load args dict and instantiate ModelArgs
        saved_args_data = checkpoint.get('model_args', checkpoint.get('model_args_dict')) # Check both keys
        if not saved_args_data: exit("Error: model_args not found in checkpoint.")
        if not isinstance(saved_args_data, dict): saved_args_dict = dataclasses.asdict(saved_args_data)
        else: saved_args_dict = saved_args_data
        init_field_names = {f.name for f in dataclasses.fields(ModelArgs) if f.init}
        filtered_args_dict = {k: v for k, v in saved_args_dict.items() if k in init_field_names}
        config_inf = ModelArgs(**filtered_args_dict)

        print(f"Loaded model config from checkpoint: {config_inf}")

        # --- Instantiate Model ---
        model_inf = Llama(config_inf) # Instantiate on CPU
        print("Model instantiated on CPU.")

        # --- Load Weights ---
        model_inf.load_state_dict(checkpoint['model_state_dict'])
        print("Model weights loaded.")
        model_inf.to(device) # Move model to target device
        print(f"Model moved to {device}.")

        # --- Prepare for Inference ---
        model_inf.eval()
        if device.type == 'cuda':
            try: model_inf = model_inf.half(); print("Converted loaded model to float16 for inference.")
            except Exception as e: print(f"Could not convert model to float16: {e}")

    except Exception as e: exit(f"Error loading checkpoint or instantiating model: {e}")

    # --- Load Tokenizer ---
    tokenizer_name_inf = "deepseek-ai/DeepSeek-R1"
    print(f"Loading tokenizer: {tokenizer_name_inf}")
    try:
        tokenizer_inf = AutoTokenizer.from_pretrained(tokenizer_name_inf, trust_remote_code=True)
        if tokenizer_inf.pad_token is None:
            if tokenizer_inf.eos_token: tokenizer_inf.pad_token = tokenizer_inf.eos_token
            else: tokenizer_inf.add_special_tokens({'pad_token': '[PAD]'})
        print("Tokenizer loaded.")
    except Exception as e: exit(f"Error loading tokenizer: {e}")

    # --- Run Generation ---
    print(f"\n--- Running Generation with Loaded Checkpoint ({os.path.basename(checkpoint_path)}) ---") # Updated print
    prompt_inf = "Valkyria Chronicles is a tactical role-playing game developed and published by"
    max_gen_len = 100
    gen_temperature = 0.7
    gen_top_k = 50
    gen_top_p = 0.9

    try:
        start_time_inf = time.time()
        _ = generate(
            model=model_inf, tokenizer=tokenizer_inf, prompt=prompt_inf,
            max_new_tokens=max_gen_len, temperature=gen_temperature,
            top_k=gen_top_k, top_p=gen_top_p
        )
        end_time_inf = time.time()
        print(f"\nInference duration: {datetime.timedelta(seconds=int(end_time_inf - start_time_inf))}")
        print("\n(Output quality depends heavily on limited training. Expect limited coherence.)")
    except Exception as e: print(f"\nAn error occurred during generation: {e}"); traceback.print_exc()

    print("\n--- Inference Script Section Finished ---")