Blame for model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import json
import os
from dataclasses import dataclass, asdict

VOCAB_SIZE = 50257
BLOCK_SIZE = 1024
LAYERS = 12
EMBEDDING_DIM = 768
NUM_HEADS = 12
DROPOUT = 0.1


@dataclass
class GPTConfig:
    vocab_size: int = VOCAB_SIZE
    block_size: int = BLOCK_SIZE
    n_layers: int = LAYERS
    n_embd: int = EMBEDDING_DIM
    n_heads: int = NUM_HEADS
    dropout: float = DROPOUT


class CausalSelfAttention(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        if config.n_embd % config.n_heads != 0:
            raise ValueError("n_embd must be divisible by n_heads")

        self.n_heads = config.n_heads
        self.head_dim = config.n_embd // config.n_heads
        self.dropout = config.dropout

        self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.out_proj = nn.Linear(config.n_embd, config.n_embd)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        causal_mask = torch.tril(torch.ones(config.block_size, config.block_size, dtype=torch.bool))
        self.register_buffer("causal_mask", causal_mask, persistent=False)

    def forward(self, x, attention_mask=None):
        batch_size, seq_len, channels = x.size()
        qkv = self.qkv(x)
        q, k, v = qkv.split(channels, dim=2)

        q = q.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)

        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        causal = self.causal_mask[:seq_len, :seq_len].view(1, 1, seq_len, seq_len)

        if attention_mask is not None:
            key_mask = attention_mask[:, None, None, :].bool()
            combined_mask = causal & key_mask
        else:
            combined_mask = causal

        # Use finite mask values to avoid NaNs on fully masked rows, then zero masked probs.
        mask_value = torch.finfo(att.dtype).min
        att = att.masked_fill(~combined_mask, mask_value)
        att = F.softmax(att, dim=-1)
        att = att.masked_fill(~combined_mask, 0.0)
        att = self.attn_dropout(att)

        y = att @ v
        y = y.transpose(1, 2).contiguous().view(batch_size, seq_len, channels)
        y = self.resid_dropout(self.out_proj(y))
        return y


class MLP(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        self.fc_in = nn.Linear(config.n_embd, 4 * config.n_embd)
        self.fc_out = nn.Linear(4 * config.n_embd, config.n_embd)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        x = self.fc_in(x)
        x = F.gelu(x)
        x = self.fc_out(x)
        return self.dropout(x)


class TransformerBlock(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        self.ln_1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x, attention_mask=None):
        x = x + self.attn(self.ln_1(x), attention_mask=attention_mask)
        x = x + self.mlp(self.ln_2(x))
        return x


class GPTModel(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()
        self.config = config
        self.token_emb = nn.Embedding(config.vocab_size, config.n_embd)
        self.pos_emb = nn.Embedding(config.block_size, config.n_embd)
        self.drop = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
        self.ln_f = nn.LayerNorm(config.n_embd)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        self.lm_head.weight = self.token_emb.weight
        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def forward(self, input_ids, attention_mask=None, labels=None):
        batch_size, seq_len = input_ids.shape
        if seq_len > self.config.block_size:
            raise ValueError(f"Sequence length {seq_len} exceeds block size {self.config.block_size}")

        pos = torch.arange(0, seq_len, device=input_ids.device).unsqueeze(0)
        x = self.token_emb(input_ids) + self.pos_emb(pos)
        x = self.drop(x)

        for block in self.blocks:
            x = block(x, attention_mask=attention_mask)

        x = self.ln_f(x)
        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                labels.view(-1),
                ignore_index=-100,
            )
        return logits, loss

    @torch.no_grad()
    def generate(self, input_ids, attention_mask=None, max_new_tokens=50, temperature=1.0, top_k=None, top_p=0.95, eos_token_id=None):
        self.eval()
        for _ in range(max_new_tokens):
            idx_cond = input_ids[:, -self.config.block_size :]
            if attention_mask is not None:
                mask_cond = attention_mask[:, -self.config.block_size :]
            else:
                mask_cond = None

            logits, _ = self(idx_cond, attention_mask=mask_cond)
            logits = logits[:, -1, :] / max(temperature, 1e-6)

            # Use one filtering strategy at a time to avoid ambiguous distributions.
            if top_p is not None and 0 < top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                probs = F.softmax(sorted_logits, dim=-1)
                cumulative = torch.cumsum(probs, dim=-1)
                sorted_mask = cumulative > top_p
                sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
                sorted_mask[..., 0] = False
                sorted_logits = sorted_logits.masked_fill(sorted_mask, -float("inf"))
                logits = torch.full_like(logits, -float("inf"))
                logits.scatter_(1, sorted_indices, sorted_logits)
            elif top_k is not None and top_k > 0:
                values, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < values[:, [-1]]] = -float("inf")

            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            input_ids = torch.cat((input_ids, next_token), dim=1)
            if attention_mask is not None:
                attention_mask = torch.cat(
                    (attention_mask, torch.ones_like(next_token, device=attention_mask.device)),
                    dim=1,
                )

            if eos_token_id is not None and torch.all(next_token == eos_token_id):
                break

        return input_ids


class miniGPT:
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

        if self.tokenizer.pad_token_id is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

    @classmethod
    def from_config(cls, config: GPTConfig, tokenizer):
        model = GPTModel(config)
        return cls(model=model, tokenizer=tokenizer)

    @classmethod
    def from_tokenizer_name(cls, tokenizer_name="gpt2", config=None):
        from transformers import AutoTokenizer

        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
        if tokenizer.pad_token_id is None:
            tokenizer.pad_token = tokenizer.eos_token

        if config is None:
            config = GPTConfig(vocab_size=len(tokenizer))
        else:
            config.vocab_size = len(tokenizer)

        return cls.from_config(config=config, tokenizer=tokenizer)

    def to(self, device):
        self.model = self.model.to(device)
        return self

    def train(
        self,
        data,
        epochs=10,
        batch_size=32,
        learning_rate=3e-4,
        weight_decay=0.1,
        grad_clip=1.0,
        warmup_ratio=0.03,
        min_lr_ratio=0.1,
    ):
        device = next(self.model.parameters()).device
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=learning_rate, weight_decay=weight_decay)
        pad_token_id = self.tokenizer.pad_token_id

        self.model.train()

        steps_per_epoch = max(1, (len(data) + batch_size - 1) // batch_size)
        total_steps = max(1, epochs * steps_per_epoch)
        warmup_steps = int(total_steps * warmup_ratio)
        global_step = 0

        def _set_lr(step):
            if warmup_steps > 0 and step < warmup_steps:
                scale = step / max(1, warmup_steps)
            else:
                decay_steps = max(1, total_steps - warmup_steps)
                progress = (step - warmup_steps) / decay_steps
                progress = min(max(progress, 0.0), 1.0)
                cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
                scale = min_lr_ratio + (1.0 - min_lr_ratio) * cosine

            current_lr = learning_rate * scale
            for group in optimizer.param_groups:
                group["lr"] = current_lr
            return current_lr

        for epoch in range(epochs):
            epoch_loss = 0.0
            epoch_batches = 0

            for i in range(0, len(data), batch_size):
                batch = data[i : i + batch_size]
                inputs = self.tokenizer(
                    batch,
                    return_tensors="pt",
                    padding=True,
                    truncation=True,
                    max_length=self.model.config.block_size,
                )
                input_ids = inputs["input_ids"].to(device)
                attention_mask = inputs["attention_mask"].to(device)

                labels = input_ids.clone()
                labels[:, :-1] = input_ids[:, 1:]
                labels[:, -1] = -100
                labels = labels.masked_fill(input_ids == pad_token_id, -100)

                _, loss = self.model(input_ids, attention_mask=attention_mask, labels=labels)

                current_lr = _set_lr(global_step)
                optimizer.zero_grad(set_to_none=True)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), grad_clip)
                optimizer.step()
                global_step += 1

                loss_value = loss.item()
                epoch_loss += loss_value
                epoch_batches += 1

                avg_loss = epoch_loss / epoch_batches
                print(
                    f"Epoch {epoch + 1}/{epochs}, Batch {i // batch_size + 1}/{steps_per_epoch}, "
                    f"Loss: {avg_loss:.4f}, LR: {current_lr:.2e}"
                )

            if epoch_batches > 0:
                print(f"Epoch {epoch + 1} average loss: {epoch_loss / epoch_batches:.4f}")

    @torch.no_grad()
    def generate(self, prompt, max_length=50, temperature=0.9, top_k=None, top_p=0.95):
        device = next(self.model.parameters()).device
        self.model.eval()

        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=self.model.config.block_size)
        input_ids = inputs["input_ids"].to(device)
        attention_mask = inputs["attention_mask"].to(device)

        output_ids = self.model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_new_tokens=max_length,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            eos_token_id=self.tokenizer.eos_token_id,
        )
        return self.tokenizer.decode(output_ids[0], skip_special_tokens=True)

    def save(self, output_dir):
        os.makedirs(output_dir, exist_ok=True)
        with open(os.path.join(output_dir, "config.json"), "w", encoding="utf-8") as config_file:
            json.dump(asdict(self.model.config), config_file, indent=2)
        torch.save(self.model.state_dict(), os.path.join(output_dir, "weights.pt"))
        self.tokenizer.save_pretrained(output_dir)

    @classmethod
    def load(cls, model_dir, map_location=None):
        from transformers import AutoTokenizer

        config_path = os.path.join(model_dir, "config.json")
        weights_path = os.path.join(model_dir, "weights.pt")

        if os.path.exists(config_path) and os.path.exists(weights_path):
            with open(config_path, "r", encoding="utf-8") as config_file:
                config_dict = json.load(config_file)
                # Map HuggingFace field names to our GPTConfig field names
                filtered_dict = {}
                if 'vocab_size' in config_dict:
                    filtered_dict['vocab_size'] = config_dict['vocab_size']
                if 'n_ctx' in config_dict:
                    filtered_dict['block_size'] = config_dict['n_ctx']
                if 'n_layer' in config_dict:
                    filtered_dict['n_layers'] = config_dict['n_layer']
                if 'n_embd' in config_dict:
                    filtered_dict['n_embd'] = config_dict['n_embd']
                if 'n_head' in config_dict:
                    filtered_dict['n_heads'] = config_dict['n_head']
                # Use any available dropout field
                if 'resid_pdrop' in config_dict:
                    filtered_dict['dropout'] = config_dict['resid_pdrop']
                elif 'attn_pdrop' in config_dict:
                    filtered_dict['dropout'] = config_dict['attn_pdrop']
                elif 'embd_pdrop' in config_dict:
                    filtered_dict['dropout'] = config_dict['embd_pdrop']
            
            # Load weights - check if it's a full checkpoint or just state_dict
            loaded = torch.load(weights_path, map_location=map_location)
            if isinstance(loaded, dict) and "state_dict" in loaded:
                state_dict = loaded["state_dict"]
            else:
                state_dict = loaded
            
            # Infer actual dimensions from state_dict to handle config/weight mismatches
            if "token_emb.weight" in state_dict:
                actual_embd = state_dict["token_emb.weight"].shape[1]
                filtered_dict['n_embd'] = actual_embd
            
            if "pos_emb.weight" in state_dict:
                actual_block_size = state_dict["pos_emb.weight"].shape[0]
                filtered_dict['block_size'] = actual_block_size
            
            # Count actual layers from state_dict
            layer_count = 0
            for key in state_dict.keys():
                if key.startswith("blocks."):
                    layer_num = int(key.split(".")[1])
                    layer_count = max(layer_count, layer_num + 1)
            if layer_count > 0:
                filtered_dict['n_layers'] = layer_count
            
            config = GPTConfig(**filtered_dict)
        else:
            checkpoint = torch.load(os.path.join(model_dir, "model.pt"), map_location=map_location)
            config = GPTConfig(**checkpoint["config"])
            state_dict = checkpoint["state_dict"]

        model = GPTModel(config)
        model.load_state_dict(state_dict, strict=False)

        tokenizer = AutoTokenizer.from_pretrained(model_dir)
        if tokenizer.pad_token_id is None:
            tokenizer.pad_token = tokenizer.eos_token

        return cls(model=model, tokenizer=tokenizer)