files for inference

Architecture, Huggingface generator and interaction code

Architecture.py

@@ -0,0 +1,146 @@
+import torch
+import torch.nn as nn
+import math
+
+
+class MaskedMultiHeadedSelfAttention(nn.Module):
+    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
+        super().__init__()
+        assert (d_out % num_heads == 0), "d_out must be divisible by num_heads"
+
+        self.d_out     = d_out
+        self.num_heads = num_heads
+        self.head_dim  = d_out // num_heads
+
+        # ── Fused QKV projection (one matmul instead of three) ──────────────
+        self.W_qkv   = nn.Linear(d_in, 3 * d_out, bias=qkv_bias)
+        self.out_proj = nn.Linear(d_out, d_out, bias=False)
+        self.dropout  = nn.Dropout(dropout)
+
+    def forward(self, x):
+        batch_size, num_tokens, d_in = x.shape
+
+        # ── Single fused projection then split ──────────────────────────────
+        qkv     = self.W_qkv(x)                          # (B, T, 3 * d_out)
+        q, k, v = qkv.split(self.d_out, dim=-1)          # each: (B, T, d_out)
+
+        # Reshape to (B, num_heads, T, head_dim)
+        q = q.view(batch_size, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch_size, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch_size, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
+
+        # ── Flash Attention — no manual mask, is_causal handles it ──────────
+        context_vec = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=None,
+            dropout_p=self.dropout.p if self.training else 0.0,
+            is_causal=True
+        )
+
+        # Merge heads and project
+        context_vec = context_vec.transpose(1, 2).reshape(batch_size, num_tokens, self.d_out)
+        context_vec = self.out_proj(context_vec)
+        return context_vec
+
+
+class FeedForward(nn.Module):
+    def __init__(self, configuration):
+        super().__init__()
+        dim = configuration["embedding_dim"]
+        self.layers = nn.Sequential(
+            nn.Linear(dim, 4 * dim),
+            nn.GELU(approximate='tanh'),   # fused kernel, faster than manual tanh approx
+            nn.Linear(4 * dim, dim),
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, configuration):
+        super().__init__()
+        self.attention    = MaskedMultiHeadedSelfAttention(
+            d_in=configuration["embedding_dim"],
+            d_out=configuration["embedding_dim"],
+            context_length=configuration["context_length"],
+            num_heads=configuration["n_heads"],
+            dropout=configuration["dropout_rate"],
+            qkv_bias=configuration["qkv_bias"]
+        )
+        self.feed_forward = FeedForward(configuration)
+
+        # ── nn.LayerNorm uses a fused CUDA kernel — faster than custom impl ──
+        self.norm1        = nn.LayerNorm(configuration["embedding_dim"])
+        self.norm2        = nn.LayerNorm(configuration["embedding_dim"])
+        self.drop_shortcut = nn.Dropout(configuration["dropout_rate"])
+
+    def forward(self, x):
+        # ── Attention block with residual ────────────────────────────────────
+        shortcut = x
+        x = self.norm1(x)
+        x = self.attention(x)
+        x = self.drop_shortcut(x)
+        x = x + shortcut
+
+        # ── Feed-forward block with residual ─────────────────────────────────
+        shortcut = x
+        x = self.norm2(x)
+        x = self.feed_forward(x)
+        x = self.drop_shortcut(x)
+        x = x + shortcut
+
+        return x
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, configuration):
+        super().__init__()
+        self.config = configuration
+
+        self.token_embedding = nn.Embedding(configuration["vocab_size"], configuration["embedding_dim"])
+        self.pos_embedding   = nn.Embedding(configuration["context_length"], configuration["embedding_dim"])
+        self.drop_embedding  = nn.Dropout(configuration["dropout_rate"])
+
+        self.transformer_blocks = nn.Sequential(
+            *[TransformerBlock(configuration) for _ in range(configuration["n_layers"])]
+        )
+
+        # ── Final norm also switched to nn.LayerNorm ─────────────────────────
+        self.final_norm = nn.LayerNorm(configuration["embedding_dim"])
+        self.out_head   = nn.Linear(configuration["embedding_dim"], configuration["vocab_size"], bias=False)
+
+        # Weight tying — output projection shares weights with token embedding
+        self.out_head.weight = self.token_embedding.weight
+
+        # ── GPT-2 style weight initialization ────────────────────────────────
+        self.apply(self._init_weights)
+
+        # Scale residual-path projections by 1/√(2*n_layers) to prevent the
+        # residual stream from blowing up in deep networks. Targets:
+        #   - attention out_proj  (ends with 'out_proj.weight')
+        #   - FFN second linear   (ends with 'layers.2.weight')
+        residual_std = 0.02 / math.sqrt(2 * configuration["n_layers"])
+        for pn, p in self.named_parameters():
+            if pn.endswith('out_proj.weight') or pn.endswith('layers.2.weight'):
+                torch.nn.init.normal_(p, mean=0.0, std=residual_std)
+
+    def _init_weights(self, module):
+        """Initialize weights following GPT-2: N(0, 0.02) for all projections."""
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, in_idx):
+        batch_size, seq_len = in_idx.shape
+        tok_embeds = self.token_embedding(in_idx)
+        pos_embeds = self.pos_embedding(torch.arange(seq_len, device=in_idx.device))
+        x = tok_embeds + pos_embeds
+        x = self.drop_embedding(x)
+        x = self.transformer_blocks(x)
+        x = self.final_norm(x)
+        logits = self.out_head(x)
+        return logits

GeneratorHF.py

@@ -0,0 +1,199 @@
+"""
+HuggingFace-compatible generator for custom LanguageModel.
+
+Wraps the custom model in a minimal HF-compatible interface so we can use
+transformers.generate() with all its bells and whistles (beam search,
+contrastive search, repetition penalty, etc.) while keeping our own weights.
+
+Usage:
+    from GeneratorHF import HFTextGenerator
+
+    gen = HFTextGenerator(model, tokenizer, device, context_size=1024)
+    text = gen.generate("Once upon a time", max_new_tokens=200)
+"""
+
+import torch
+import torch.nn as nn
+from transformers import GenerationConfig, GenerationMixin, PretrainedConfig
+
+
+class LanguageModelHF(nn.Module, GenerationMixin):
+    """Thin wrapper that makes our custom LanguageModel compatible with
+    HuggingFace's GenerationMixin (generate, beam_search, sample, etc.)."""
+
+    _is_stateful = False
+    _supports_cache_class = False
+
+    def __init__(self, model, context_size, device):
+        super().__init__()
+        self.model = model
+        self.config = _make_hf_config(model, context_size)
+        self.generation_config = GenerationConfig(
+            max_new_tokens=200,
+            do_sample=True,
+            temperature=0.8,
+            top_k=40,
+            top_p=0.95,
+            repetition_penalty=1.1,
+            eos_token_id=50256,
+            pad_token_id=50256,
+        )
+        self.main_input_name = "input_ids"
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def forward(self, input_ids, attention_mask=None, **kwargs):
+        # Trim to context window (HF might feed longer sequences)
+        input_ids = input_ids[:, -self.config.n_positions:]
+        logits = self.model(input_ids)
+        return _CausalLMOutput(logits=logits)
+
+    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
+        return {"input_ids": input_ids}
+
+    def can_generate(self):
+        return True
+
+    def _reorder_cache(self, past, beam_idx):
+        return past
+
+
+class _CausalLMOutput:
+    """Minimal output container matching HF's CausalLMOutput interface."""
+    def __init__(self, logits):
+        self.logits = logits
+        self.past_key_values = None
+        self.hidden_states = None
+        self.attentions = None
+
+    def __getitem__(self, key):
+        return getattr(self, key)
+
+    def __contains__(self, key):
+        return hasattr(self, key) and getattr(self, key) is not None
+
+    def keys(self):
+        return [k for k in ["logits", "past_key_values", "hidden_states", "attentions"]
+                if getattr(self, k, None) is not None]
+
+
+class _MinimalConfig(PretrainedConfig):
+    """Config that satisfies HF GenerationMixin by inheriting PretrainedConfig."""
+    model_type = "custom_gpt2"
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.is_encoder_decoder = False
+        self.n_positions = kwargs.get("n_positions", 1024)
+        self.vocab_size = kwargs.get("vocab_size", 50257)
+        self.eos_token_id = 50256
+        self.pad_token_id = 50256
+        self.bos_token_id = 50256
+        self.use_cache = False
+        self.num_hidden_layers = kwargs.get("num_hidden_layers", 12)
+
+
+def _make_hf_config(model, context_size):
+    vocab_size = 50257
+    n_positions = context_size
+    num_hidden_layers = 12
+    if hasattr(model, 'config') and isinstance(model.config, dict):
+        vocab_size = model.config.get("vocab_size", 50257)
+        n_positions = model.config.get("context_length", context_size)
+        num_hidden_layers = model.config.get("n_layers", 12)
+
+    cfg = _MinimalConfig(vocab_size=vocab_size, n_positions=n_positions, num_hidden_layers=num_hidden_layers)
+    return cfg
+
+
+class HFTextGenerator:
+    """Drop-in replacement for TextGenerator using HF generate()."""
+
+    def __init__(self, model, tokenizer, device, context_size=1024):
+        self.tokenizer = tokenizer
+        self.device = device
+        self.context_size = context_size
+        self.hf_model = LanguageModelHF(model, context_size, device)
+        self.hf_model.to(device)
+        self.hf_model.eval()
+
+    def generate(self, prompt, max_new_tokens=200,
+                 temperature=0.8, top_k=40, top_p=0.95,
+                 repetition_penalty=1.1, num_beams=1,
+                 do_sample=True, eos_id=50256):
+        """Generate text from a string prompt. Returns the generated string."""
+        input_ids = torch.tensor(
+            [self.tokenizer.encode(prompt, allowed_special={"<|endoftext|>"})]
+        ).to(self.device)
+
+        gen_config = GenerationConfig(
+            max_new_tokens=max_new_tokens,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            num_beams=num_beams,
+            eos_token_id=eos_id,
+            pad_token_id=eos_id,
+        )
+
+        attention_mask = torch.ones_like(input_ids)
+
+        with torch.no_grad():
+            output_ids = self.hf_model.generate(
+                input_ids,
+                attention_mask=attention_mask,
+                generation_config=gen_config,
+            )
+
+        # Decode only the newly generated tokens
+        new_tokens = output_ids[0, input_ids.shape[1]:]
+        
+        # Filter out the EOS token if it was generated
+        if len(new_tokens) > 0 and new_tokens[-1].item() == eos_id:
+            new_tokens = new_tokens[:-1]
+            
+        decoded_text = self.tokenizer.decode(new_tokens.tolist())
+        return decoded_text.replace("<|endoftext|>", "").strip()
+
+    def generate_ids(self, idx, max_new_tokens=200,
+                     temperature=0.8, top_k=40, top_p=0.95,
+                     repetition_penalty=1.1, num_beams=1,
+                     do_sample=True, eos_id=50256):
+        """Generate from token IDs tensor (same interface as original).
+        Returns full sequence including prompt."""
+        gen_config = GenerationConfig(
+            max_new_tokens=max_new_tokens,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            num_beams=num_beams,
+            eos_token_id=eos_id,
+            pad_token_id=eos_id,
+        )
+
+        idx = idx.to(self.device)
+        attention_mask = torch.ones_like(idx)
+
+        with torch.no_grad():
+            output_ids = self.hf_model.generate(
+                idx,
+                attention_mask=attention_mask,
+                generation_config=gen_config,
+            )
+
+        return output_ids
+
+    # ── Backward-compatible helpers (match TextGenerator interface) ───────
+    def token_ids_to_text(self, token_ids):
+        flat = token_ids.squeeze(0)
+        return self.tokenizer.decode(flat.tolist())
+
+    def text_to_token_ids(self, text):
+        encoded = self.tokenizer.encode(text, allowed_special={"<|endoftext|>"})
+        return torch.tensor(encoded).unsqueeze(0).to(self.device)

chat.py

@@ -0,0 +1,131 @@
+"""
+Interactive chat with your GPT-2 model — Notebook version
+Run Cell 1 and Cell 2 in order, then use the chat widget.
+You can copy the code into colab cells. ensure you rename to your model path
+"""
+
+# ── Cell 1: Setup & load model ────────────────────────────────────────────────
+import torch
+import tiktoken
+import ipywidgets as widgets
+from IPython.display import display, HTML
+from Architecture import LanguageModel
+from GeneratorHF import HFTextGenerator
+
+CHECKPOINT_PATH = "/content/drive/MyDrive/GPT2-small/refinedGPT2small.pth" #on google colab
+DEVICE          = "cuda" if torch.cuda.is_available() else "cpu"
+
+def load_model(checkpoint_path, device):
+    print(f"Loading model from {checkpoint_path}...")
+    ckpt       = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    config     = ckpt["config"]
+    model      = LanguageModel(config)
+    state_dict = ckpt["model_state_dict"]
+
+    cleaned = {}
+    for k, v in state_dict.items():
+        for prefix in ["module.", "_orig_mod."]:
+            while k.startswith(prefix):
+                k = k[len(prefix):]
+        cleaned[k] = v
+
+    model.load_state_dict(cleaned)
+    model.to(device)
+    model.eval()
+
+    params   = sum(p.numel() for p in model.parameters()) / 1e6
+    val_loss = ckpt.get("loss", "?")
+    print(f"  {params:.1f}M params | val_loss {val_loss} | device {device}")
+    return model, config
+
+model, config = load_model(CHECKPOINT_PATH, DEVICE)
+enc           = tiktoken.get_encoding("gpt2")
+gen           = HFTextGenerator(model, enc, DEVICE, context_size=config["context_length"])
+print("Ready.")
+
+
+# ── Cell 2: Chat widget ───────────────────────────────────────────────────────
+
+# ── Sampling controls ─────────────────────────────────────────────────────────
+temp_slider = widgets.FloatSlider(
+    value=0.8, min=0.1, max=2.0, step=0.05,
+    description="Temperature",
+    style={"description_width": "90px"},
+    layout=widgets.Layout(width="400px")
+)
+topk_slider = widgets.IntSlider(
+    value=40, min=1, max=200, step=1,
+    description="Top-k",
+    style={"description_width": "90px"},
+    layout=widgets.Layout(width="400px")
+)
+tokens_slider = widgets.IntSlider(
+    value=150, min=10, max=500, step=10,
+    description="Max tokens",
+    style={"description_width": "90px"},
+    layout=widgets.Layout(width="400px")
+)
+
+# ── Input & buttons ───────────────────────────────────────────────────────────
+prompt_box = widgets.Textarea(
+    placeholder="Type your prompt here...",
+    layout=widgets.Layout(width="700px", height="70px")
+)
+generate_btn = widgets.Button(
+    description="Generate",
+    button_style="primary",
+    layout=widgets.Layout(width="120px")
+)
+clear_btn = widgets.Button(
+    description="Clear",
+    button_style="warning",
+    layout=widgets.Layout(width="80px")
+)
+
+# ── Output area ───────────────────────────────────────────────────────────────
+output_area = widgets.Output()
+
+def on_generate(b):
+    prompt = prompt_box.value.strip()
+    if not prompt:
+        return
+
+    with output_area:
+        print(f"\nYOU   > {prompt}")
+        print("MODEL > ", end="", flush=True)
+
+        continuation = gen.generate(
+            prompt,
+            max_new_tokens=tokens_slider.value,
+            temperature=temp_slider.value,
+            top_k=topk_slider.value,
+        )
+        print(continuation)
+        print("-" * 70)
+
+def on_clear(b):
+    output_area.clear_output()
+    prompt_box.value = ""
+
+generate_btn.on_click(on_generate)
+clear_btn.on_click(on_clear)
+
+# ── Layout ────────────────────────────────────────────────────────────────────
+controls = widgets.VBox([
+    widgets.HTML("<b>Sampling settings</b>"),
+    temp_slider,
+    topk_slider,
+    tokens_slider,
+])
+buttons = widgets.HBox([generate_btn, clear_btn])
+ui      = widgets.VBox([
+    widgets.HTML("<h3>GPT-2 Small — Chat</h3>"),
+    controls,
+    widgets.HTML("<br><b>Prompt</b>"),
+    prompt_box,
+    buttons,
+    widgets.HTML("<br><b>Output</b>"),
+    output_area,
+])
+
+display(ui)