Model-Training(Pretraining)/inference.py

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()