import os import math import time import inspect from dataclasses import dataclass import torch import torch.nn as nn from torch.nn import functional as F import tiktoken # --------- RMSNormalization ----- # --- RMSNorm Class --- 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): output = self._norm(x.float()).type_as(x) return output * self.weight # --- RoPE and KV Cache Components --- def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0): freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) t = torch.arange(end, device=freqs.device) freqs = torch.outer(t, freqs).float() freqs_cis = torch.polar(torch.ones_like(freqs), freqs) return freqs_cis # Rotary positional encoding ! def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor): xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) freqs_cis = freqs_cis.unsqueeze(0).unsqueeze(1) xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) return xq_out.type_as(xq), xk_out.type_as(xk) # --- CausalSelfAttention with GQA --- class CausalSelfAttention(nn.Module): def __init__(self, config): super().__init__() assert config.n_embd % config.n_head == 0 self.n_head = config.n_head self.n_kv_heads = config.n_kv_heads self.n_embd = config.n_embd self.head_dim = self.n_embd // self.n_head assert self.n_head % self.n_kv_heads == 0, "n_head must be divisible by n_kv_heads" self.q_proj = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False) self.k_proj = nn.Linear(self.n_embd, self.n_kv_heads * self.head_dim, bias=False) self.v_proj = nn.Linear(self.n_embd, self.n_kv_heads * self.head_dim, bias=False) self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False) self.register_buffer("freqs_cis", precompute_freqs_cis(self.head_dim, config.block_size * 2)) def forward(self, x, kv_cache=None, start_pos=0): B, T, C = x.size() q = self.q_proj(x).view(B, T, self.n_head, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2) q, k = apply_rotary_emb(q, k, freqs_cis=self.freqs_cis[start_pos : start_pos + T]) if kv_cache is not None: cache_k, cache_v = kv_cache k = torch.cat([cache_k, k], dim=2) v = torch.cat([cache_v, v], dim=2) updated_kv_cache = (k, v) n_repeats = self.n_head // self.n_kv_heads if n_repeats > 1: k = k.repeat_interleave(n_repeats, dim=1) v = v.repeat_interleave(n_repeats, dim=1) y = F.scaled_dot_product_attention(q, k, v, is_causal=(kv_cache is None)) y = y.transpose(1, 2).contiguous().view(B, T, C) y = self.c_proj(y) return y, updated_kv_cache # --- MLP with SwiGLU --- class MLP(nn.Module): def __init__(self, config): super().__init__() hidden_dim = 4 * config.n_embd hidden_dim = int(2 * hidden_dim / 3) self.w1 = nn.Linear(config.n_embd, hidden_dim, bias=False) self.w3 = nn.Linear(config.n_embd, hidden_dim, bias=False) self.w2 = nn.Linear(hidden_dim, config.n_embd, bias=False) def forward(self, x): return self.w2(F.silu(self.w1(x)) * self.w3(x)) # --- Block with RMSNorm --- class Block(nn.Module): def __init__(self, config): super().__init__() self.ln_1 = RMSNorm(config.n_embd) self.attn = CausalSelfAttention(config) self.ln_2 = RMSNorm(config.n_embd) self.mlp = MLP(config) def forward(self, x, kv_cache=None, start_pos=0): attn_output, updated_kv_cache = self.attn(self.ln_1(x), kv_cache=kv_cache, start_pos=start_pos) x = x + attn_output x = x + self.mlp(self.ln_2(x)) return x, updated_kv_cache # --- GPTConfig with GQA parameters --- @dataclass class GPTConfig: block_size: int = 1024 vocab_size: int = 50257 n_layer: int = 24 n_head: int = 16 n_kv_heads: int = 8 n_embd: int = 1024 class GPT(nn.Module): def __init__(self, config): super().__init__() self.config = config self.transformer = nn.ModuleDict(dict( wte = nn.Embedding(config.vocab_size, config.n_embd), h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), ln_f = RMSNorm(config.n_embd), )) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) self.transformer.wte.weight = self.lm_head.weight def forward(self, idx, kv_cache=None, start_pos=0): B, T = idx.size() assert T <= self.config.block_size, f"Sequence length {T} exceeds block size {self.config.block_size}" x = self.transformer.wte(idx) if kv_cache is None: kv_cache = [None] * self.config.n_layer updated_kv_cache = [] for i, block in enumerate(self.transformer.h): x, new_cache = block(x, kv_cache=kv_cache[i], start_pos=start_pos) updated_kv_cache.append(new_cache) x = self.transformer.ln_f(x) logits = self.lm_head(x) return logits, updated_kv_cache def main(): # --- Configuration ---clea model_path = './model_472000.pt' # Path to your model weights device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {device}") device_type = "cuda" if device.startswith("cuda") else "cpu" torch.manual_seed(1337) if torch.cuda.is_available(): torch.cuda.manual_seed(1337) # --- Model Loading --- print("Loading model...") if 'weights_only' in inspect.signature(torch.load).parameters: checkpoint = torch.load(model_path, map_location=device, weights_only=False) else: checkpoint = torch.load(model_path, map_location=device) config = checkpoint['config'] model = GPT(config) state_dict = checkpoint['model'] unwanted_prefix = '_orig_mod.' for k,v in list(state_dict.items()): if k.startswith(unwanted_prefix): state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k) model.load_state_dict(state_dict) model.to(device) model.eval() print("Model loaded successfully.") prompt = input("what's up: ") num_return_sequences = int(input("times - response: ")) # max token limit`` max_length = 750 # --- Tokenizer --- enc = tiktoken.get_encoding("gpt2") # --- Generation Loop for multiple sequences --- print("\nStarting generation...") for seq_idx in range(num_return_sequences): print(f"\n--- Generating Sample {seq_idx+1} ---") tokens = enc.encode(prompt) tokens = torch.tensor(tokens, dtype=torch.long, device=device).unsqueeze(0) xgen = tokens # This will accumulate the generated sequence for this sample sample_rng = torch.Generator(device=device) sample_rng.manual_seed(42 + seq_idx) kv_cache = None start_pos = 0 with torch.no_grad(): # Process the initial prompt with torch.autocast(device_type=device_type, dtype=torch.bfloat16): logits, kv_cache = model(xgen, kv_cache=None, start_pos=0) start_pos = xgen.size(1) # Generate token by token, accumulating in xgen while xgen.size(1) < max_length: logits = logits[:, -1, :] probs = F.softmax(logits, dim=-1) topk_probs, topk_indices = torch.topk(probs, 50, dim=-1) ix = torch.multinomial(topk_probs, 1, generator=sample_rng) xcol = torch.gather(topk_indices, -1, ix) xgen = torch.cat((xgen, xcol), dim=1) # Append the new token with torch.autocast(device_type=device_type, dtype=torch.bfloat16): logits, kv_cache = model(xcol, kv_cache=kv_cache, start_pos=start_pos) start_pos += 1 # --- Decoding and Output for the complete sample --- full_tokens = xgen[0, :max_length].tolist() decoded_text = enc.decode(full_tokens) print(f"Sample {seq_idx+1}: {decoded_text}") print("-" * 30) if __name__ == "__main__": main()