1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 115,
5  "end_line": 140,
6  "content": "class TokenEmbedding(nn.Module):\n    \"\"\"Multi-component embedding: tied tokens + sinusoidal position with learnable\n    frequency scaling + bigram hash with per-layer adaptive injection.\"\"\"\n    def __init__(self, config):\n        super().__init__()\n        self.wte = nn.Embedding(config.vocab_size, config.n_embd)\n        self.drop = nn.Dropout(config.dropout)\n        self.block_size = config.block_size\n        self.n_embd = config.n_embd\n        self.vocab_size = config.vocab_size\n        self.n_layer = config.n_layer\n\n        # Learnable frequency scaling for sinusoidal position encoding\n        # Each dimension gets its own scaling factor, initialized to 1.0\n        self.pos_freq_scale = nn.Parameter(torch.ones(config.n_embd))\n        # Register buffer for the base frequency computation\n        div_term = torch.exp(torch.arange(0, config.n_embd, 2) * -(math.log(10000.0) / config.n_embd))\n        self.register_buffer('pos_div_term', div_term)\n\n        # Bigram hash embedding: 5x vocab for hash collision reduction\n        self.bigram_vocab_size = config.vocab_size * 5\n        self.bigram_embed = nn.Embedding(self.bigram_vocab_size, config.n_embd)\n        nn.init.zeros_(self.bigram_embed.weight)\n\n        # Per-layer adaptive scaling for bigram embedding injection\n        self.bigram_lambdas = nn.Parameter(torch.full((config.n_layer,), 0.05))\n        # Per-layer position encoding scale (starts at 1.0 for standard behavior)\n        self.pos_lambdas = nn.Parameter(torch.ones(config.n_layer))\n\n        self._cached_bigram = None\n\n    def _bigram_hash(self, idx):\n        \"\"\"Compute bigram hash indices from consecutive token pairs.\"\"\"\n        rand_int_1 = 36313\n        rand_int_2 = 27191\n        mod = self.bigram_vocab_size - 1\n        x = idx.to(torch.int32)\n        out = torch.zeros_like(x)\n        out[:, 0] = mod  # Position 0: no previous token\n        out[:, 1:] = torch.bitwise_xor(\n            rand_int_1 * x[:, 1:],\n            rand_int_2 * x[:, :-1]\n        ) % mod\n        return out.long()\n\n    def forward(self, idx):\n        b, t = idx.size()\n        tok_emb = self.wte(idx)\n\n        # Sinusoidal position encoding with learnable frequency scaling\n        pos = torch.arange(0, t, dtype=torch.float32, device=idx.device)\n        freqs = pos.unsqueeze(1) * (self.pos_div_term * self.pos_freq_scale)  # (t, n_embd/2)\n        pos_emb = torch.zeros(t, self.n_embd, device=idx.device)\n        pos_emb[:, 0::2] = freqs.sin()\n        pos_emb[:, 1::2] = freqs.cos()\n\n        # Compute and cache bigram embedding for per-layer injection\n        self._cached_bigram = self.bigram_embed(self._bigram_hash(idx))\n\n        return self.drop(tok_emb + pos_emb)\n\n    def get_value_embed(self, layer_idx):\n        \"\"\"Inject bigram embedding at every layer with per-layer adaptive scaling.\"\"\"\n        if self._cached_bigram is None or layer_idx >= self.n_layer:\n            return None\n        # Combine bigram and position scaling for this layer\n        bigram_part = self.bigram_lambdas[layer_idx] * self._cached_bigram\n        return bigram_part\n\n    def get_lm_head_weight(self):\n        return self.wte.weight\n\n    def get_num_pos_params(self):\n        # Only the frequency scaling parameters count as \"position params\"\n        return self.pos_freq_scale.numel()"
7}

1{
2  "content": "class TokenEmbedding(nn.Module):\n    \"\"\"Multi-component embedding: tied tokens + position + bigram hash with\n    per-layer adaptive injection and frequency-aware sinusoidal mixing.\"\"\"\n    def __init__(self, config):\n        super().__init__()\n        self.wte = nn.Embedding(config.vocab_size, config.n_embd)\n        self.wpe = nn.Embedding(config.block_size, config.n_embd)\n        self.drop = nn.Dropout(config.dropout)\n        self.block_size = config.block_size\n        self.n_embd = config.n_embd\n        self.vocab_size = config.vocab_size\n        self.n_layer = config.n_layer\n\n        # Bigram hash embedding: 5x vocab for hash collision reduction\n        self.bigram_vocab_size = config.vocab_size * 5\n        self.bigram_embed = nn.Embedding(self.bigram_vocab_size, config.n_embd)\n        nn.init.zeros_(self.bigram_embed.weight)\n\n        # Per-layer adaptive scaling for bigram embedding injection\n        self.bigram_lambdas = nn.Parameter(torch.full((config.n_layer,), 0.05))\n\n        # Learnable gate for mixing sinusoidal features into position embeddings\n        # This adds sinusoidal features with a learnable amplitude per dimension\n        self.pos_freq_scale = nn.Parameter(torch.ones(config.n_embd))\n        div_term = torch.exp(torch.arange(0, config.n_embd, 2) * -(math.log(10000.0) / config.n_embd))\n        self.register_buffer('pos_div_term', div_term)\n\n        self._cached_bigram = None\n\n    def _bigram_hash(self, idx):\n        \"\"\"Compute bigram hash indices from consecutive token pairs.\"\"\"\n        rand_int_1 = 36313\n        rand_int_2 = 27191\n        mod = self.bigram_vocab_size - 1\n        x = idx.to(torch.int32)\n        out = torch.zeros_like(x)\n        out[:, 0] = mod  # Position 0: no previous token\n        out[:, 1:] = torch.bitwise_xor(\n            rand_int_1 * x[:, 1:],\n            rand_int_2 * x[:, :-1]\n        ) % mod\n        return out.long()\n\n    def _compute_sinusoidal(self, idx):\n        \"\"\"Compute sinusoidal position encoding with learnable frequency scaling.\"\"\"\n        t = idx.size(1)\n        pos = torch.arange(0, t, dtype=torch.float32, device=idx.device)\n        # (t, n_embd/2)\n        freqs = pos.unsqueeze(1) * (self.pos_div_term * self.pos_freq_scale)\n        sin_part = freqs.sin()\n        cos_part = freqs.cos()\n        pos_emb = torch.zeros(t, self.n_embd, device=idx.device, dtype=torch.float32)\n        pos_emb[:, 0::2] = sin_part\n        pos_emb[:, 1::2] = cos_part\n        return pos_emb\n\n    def forward(self, idx):\n        b, t = idx.size()\n        tok_emb = self.wte(idx)\n        pos = torch.arange(0, t, dtype=torch.long, device=idx.device)\n        pos_emb = self.wpe(pos)\n\n        # Compute bigram embedding for per-layer injection\n        self._cached_bigram = self.bigram_embed(self._bigram_hash(idx))\n\n        return self.drop(tok_emb + pos_emb)\n\n    def get_value_embed(self, layer_idx):\n        \"\"\"Inject bigram embedding at every layer with per-layer adaptive scaling.\"\"\"\n        if self._cached_bigram is None or layer_idx >= self.n_layer:\n            return None\n        return self.bigram_lambdas[layer_idx] * self._cached_bigram\n\n    def get_lm_head_weight(self):\n        return self.wte.weight\n\n    def get_num_pos_params(self):\n        return self.wpe.weight.numel()",
3  "start_line": 115,
4  "end_line": 189,
5  "filename": "nanoGPT/custom_pretrain.py"
6}

1{
2  "end_line": 189,
3  "filename": "nanoGPT/custom_pretrain.py",
4  "op": "replace",
5  "start_line": 115,
6  "content": "class TokenEmbedding(nn.Module):\n    \"\"\"Token + position + bigram hash embedding with per-layer adaptive injection.\"\"\"\n    def __init__(self, config):\n        super().__init__()\n        self.wte = nn.Embedding(config.vocab_size, config.n_embd)\n        self.wpe = nn.Embedding(config.block_size, config.n_embd)\n        self.drop = nn.Dropout(config.dropout)\n        self.block_size = config.block_size\n        self.n_embd = config.n_embd\n        self.vocab_size = config.vocab_size\n        self.n_layer = config.n_layer\n        # Bigram hash embedding: 5x vocab for hash collision reduction\n        self.bigram_vocab_size = config.vocab_size * 5\n        self.bigram_embed = nn.Embedding(self.bigram_vocab_size, config.n_embd)\n        nn.init.zeros_(self.bigram_embed.weight)\n        # Per-layer learnable scaling for bigram embedding injection\n        self.bigram_lambdas = nn.Parameter(torch.full((config.n_layer,), 0.05))\n        self._cached_bigram = None\n\n    def _bigram_hash(self, idx):\n        \"\"\"Compute bigram hash indices from consecutive token pairs.\"\"\"\n        rand_int_1 = 36313\n        rand_int_2 = 27191\n        mod = self.bigram_vocab_size - 1\n        x = idx.to(torch.int32)\n        out = torch.zeros_like(x)\n        out[:, 0] = mod  # Position 0: no previous token\n        out[:, 1:] = torch.bitwise_xor(\n            rand_int_1 * x[:, 1:],\n            rand_int_2 * x[:, :-1]\n        ) % mod\n        return out.long()\n\n    def forward(self, idx):\n        b, t = idx.size()\n        tok_emb = self.wte(idx)\n        pos = torch.arange(0, t, dtype=torch.long, device=idx.device)\n        pos_emb = self.wpe(pos)\n        # Compute and cache bigram embedding for per-layer injection\n        self._cached_bigram = self.bigram_embed(self._bigram_hash(idx))\n        return self.drop(tok_emb + pos_emb)\n\n    def get_value_embed(self, layer_idx):\n        \"\"\"Inject bigram embedding at every layer with learnable scaling.\"\"\"\n        if self._cached_bigram is None or layer_idx >= self.n_layer:\n            return None\n        return self.bigram_lambdas[layer_idx] * self._cached_bigram\n\n    def get_lm_head_weight(self):\n        return self.wte.weight\n\n    def get_num_pos_params(self):\n        return self.wpe.weight.numel()"
7}

1{
2  "content": "class TokenEmbedding(nn.Module):\n    \"\"\"Token + position + bigram hash with position-aware adaptive injection.\n    \n    Key innovations:\n    1. Position-dependent bigram scaling: bigram context is more useful at early positions\n       where self-attention has less history to attend to\n    2. Per-layer adaptive injection: learned per-layer scaling for bigram residuals\n    3. Sinusoidal position enhancement: learnable frequency scaling on sinusoidal features\n       added to learned position embeddings for better extrapolation\n    \"\"\"\n    def __init__(self, config):\n        super().__init__()\n        self.wte = nn.Embedding(config.vocab_size, config.n_embd)\n        self.wpe = nn.Embedding(config.block_size, config.n_embd)\n        self.drop = nn.Dropout(config.dropout)\n        self.block_size = config.block_size\n        self.n_embd = config.n_embd\n        self.vocab_size = config.vocab_size\n        self.n_layer = config.n_layer\n        \n        # Bigram hash embedding: 5x vocab for hash collision reduction\n        self.bigram_vocab_size = config.vocab_size * 5\n        self.bigram_embed = nn.Embedding(self.bigram_vocab_size, config.n_embd)\n        nn.init.zeros_(self.bigram_embed.weight)\n        \n        # Per-layer scaling for bigram injection\n        self.bigram_lambdas = nn.Parameter(torch.full((config.n_layer,), 0.1))\n        \n        # Position-dependent gate for bigram injection\n        # Initialized to decay with position (more weight at early positions where\n        # self-attention has less history)\n        pos_decay = torch.exp(-torch.arange(config.block_size, dtype=torch.float32) / (config.block_size / 4))\n        self.pos_gate = nn.Parameter(pos_decay.clone())\n        \n        # Sinusoidal position enhancement with learnable frequency scaling\n        div_term = torch.exp(torch.arange(0, config.n_embd, 2).float() * -(math.log(10000.0) / config.n_embd))\n        self.register_buffer('pos_div_term', div_term)\n        self.pos_freq_scale = nn.Parameter(torch.ones(config.n_embd))\n        # Learnable mixing weight for sinusoidal enhancement (starts at 0)\n        self.sin_mix = nn.Parameter(torch.tensor(0.0))\n        \n        self._cached_bigram = None\n    \n    def _bigram_hash(self, idx):\n        \"\"\"Compute bigram hash indices from consecutive token pairs.\"\"\"\n        rand_int_1 = 36313\n        rand_int_2 = 27191\n        mod = self.bigram_vocab_size - 1\n        x = idx.to(torch.int32)\n        out = torch.zeros_like(x)\n        out[:, 0] = mod\n        out[:, 1:] = torch.bitwise_xor(\n            rand_int_1 * x[:, 1:],\n            rand_int_2 * x[:, :-1]\n        ) % mod\n        return out.long()\n    \n    def _compute_sinusoidal(self, t, device):\n        \"\"\"Compute sinusoidal position encoding with learnable frequency scaling.\"\"\"\n        pos = torch.arange(0, t, dtype=torch.float32, device=device)\n        freqs = pos.unsqueeze(1) * (self.pos_div_term * self.pos_freq_scale)\n        pos_emb = torch.zeros(t, self.n_embd, device=device, dtype=torch.float32)\n        pos_emb[:, 0::2] = freqs.sin()\n        pos_emb[:, 1::2] = freqs.cos()\n        return pos_emb\n    \n    def forward(self, idx):\n        b, t = idx.size()\n        tok_emb = self.wte(idx)\n        pos = torch.arange(0, t, dtype=torch.long, device=idx.device)\n        pos_emb = self.wpe(pos)\n        \n        # Add sinusoidal position enhancement\n        sin_emb = self._compute_sinusoidal(t, idx.device)\n        pos_emb = pos_emb + self.sin_mix * sin_emb\n        \n        # Compute and cache bigram embedding for per-layer injection\n        self._cached_bigram = self.bigram_embed(self._bigram_hash(idx))\n        \n        return self.drop(tok_emb + pos_emb)\n    \n    def get_value_embed(self, layer_idx):\n        \"\"\"Inject bigram embedding at every layer with position-aware scaling.\"\"\"\n        if self._cached_bigram is None or layer_idx >= self.n_layer:\n            return None\n        \n        # Position-dependent scaling: more weight at early positions\n        t = self._cached_bigram.size(1)\n        gate = self.pos_gate[:t]  # (T,)\n        scale = self.bigram_lambdas[layer_idx] * gate  # (T,)\n        \n        return scale.unsqueeze(-1).unsqueeze(0) * self._cached_bigram\n    \n    def get_lm_head_weight(self):\n        return self.wte.weight\n    \n    def get_num_pos_params(self):\n        return self.wpe.weight.numel()",
3  "start_line": 115,
4  "end_line": 167,
5  "filename": "nanoGPT/custom_pretrain.py"
6}

1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 115,
5  "end_line": 167,
6  "content": "class TokenEmbedding(nn.Module):\n    \"\"\"Token + position + bigram hash with position-aware adaptive injection.\n    \n    Key innovations:\n    1. Position-dependent bigram scaling: bigram context is more useful at early positions\n       where self-attention has less history to attend to\n    2. Per-layer adaptive injection: learned per-layer scaling for bigram residuals\n    3. Sinusoidal position enhancement: learnable frequency scaling on sinusoidal features\n       added to learned position embeddings for better extrapolation\n    \"\"\"\n    def __init__(self, config):\n        super().__init__()\n        self.wte = nn.Embedding(config.vocab_size, config.n_embd)\n        self.wpe = nn.Embedding(config.block_size, config.n_embd)\n        self.drop = nn.Dropout(config.dropout)\n        self.block_size = config.block_size\n        self.n_embd = config.n_embd\n        self.vocab_size = config.vocab_size\n        self.n_layer = config.n_layer\n        \n        # Bigram hash embedding: 5x vocab for hash collision reduction\n        self.bigram_vocab_size = config.vocab_size * 5\n        self.bigram_embed = nn.Embedding(self.bigram_vocab_size, config.n_embd)\n        nn.init.zeros_(self.bigram_embed.weight)\n        \n        # Per-layer scaling for bigram injection\n        self.bigram_lambdas = nn.Parameter(torch.full((config.n_layer,), 0.1))\n        \n        # Position-dependent gate for bigram injection\n        # Initialized to decay with position (more weight at early positions where\n        # self-attention has less history)\n        pos_decay = torch.exp(-torch.arange(config.block_size, dtype=torch.float32) / (config.block_size / 4))\n        self.pos_gate = nn.Parameter(pos_decay.clone())\n        \n        # Sinusoidal position enhancement with learnable frequency scaling\n        div_term = torch.exp(torch.arange(0, config.n_embd, 2).float() * -(math.log(10000.0) / config.n_embd))\n        self.register_buffer('pos_div_term', div_term)\n        self.pos_freq_scale = nn.Parameter(torch.ones(config.n_embd))\n        # Learnable mixing weight for sinusoidal enhancement (starts at 0)\n        self.sin_mix = nn.Parameter(torch.tensor(0.0))\n        \n        self._cached_bigram = None\n    \n    def _bigram_hash(self, idx):\n        \"\"\"Compute bigram hash indices from consecutive token pairs.\"\"\"\n        rand_int_1 = 36313\n        rand_int_2 = 27191\n        mod = self.bigram_vocab_size - 1\n        x = idx.to(torch.int32)\n        out = torch.zeros_like(x)\n        out[:, 0] = mod\n        out[:, 1:] = torch.bitwise_xor(\n            rand_int_1 * x[:, 1:],\n            rand_int_2 * x[:, :-1]\n        ) % mod\n        return out.long()\n    \n    def _compute_sinusoidal(self, t, device):\n        \"\"\"Compute sinusoidal position encoding with learnable frequency scaling.\"\"\"\n        pos = torch.arange(0, t, dtype=torch.float32, device=device)\n        freqs = pos.unsqueeze(1) * (self.pos_div_term * self.pos_freq_scale)\n        pos_emb = torch.zeros(t, self.n_embd, device=device, dtype=torch.float32)\n        pos_emb[:, 0::2] = freqs.sin()\n        pos_emb[:, 1::2] = freqs.cos()\n        return pos_emb\n    \n    def forward(self, idx):\n        b, t = idx.size()\n        tok_emb = self.wte(idx)\n        pos = torch.arange(0, t, dtype=torch.long, device=idx.device)\n        pos_emb = self.wpe(pos)\n        \n        # Add sinusoidal position enhancement\n        sin_emb = self._compute_sinusoidal(t, idx.device)\n        pos_emb = pos_emb + self.sin_mix * sin_emb\n        \n        # Compute and cache bigram embedding for per-layer injection\n        self._cached_bigram = self.bigram_embed(self._bigram_hash(idx))\n        \n        return self.drop(tok_emb + pos_emb)\n    \n    def get_value_embed(self, layer_idx):\n        \"\"\"Inject bigram embedding at every layer with position-aware scaling.\"\"\"\n        if self._cached_bigram is None or layer_idx >= self.n_layer:\n            return None\n        \n        # Position-dependent scaling: more weight at early positions\n        t = self._cached_bigram.size(1)\n        gate = self.pos_gate[:t]  # (T,)\n        scale = self.bigram_lambdas[layer_idx] * gate  # (T,)\n        \n        return scale.unsqueeze(-1).unsqueeze(0) * self._cached_bigram\n    \n    def get_lm_head_weight(self):\n        return self.wte.weight\n    \n    def get_num_pos_params(self):\n        return self.wpe.weight.numel()"
7}

1{
2  "op": "replace",
3  "filename": "nanoGPT/custom_pretrain.py",
4  "start_line": 115,
5  "end_line": 212,
6  "content": "class TokenEmbedding(nn.Module):\n    \"\"\"Token + position + bigram hash with position-aware adaptive injection.\n    \n    Key innovations:\n    1. Position-dependent bigram scaling: bigram context is more useful at early \n       positions where self-attention has less history to attend to\n    2. Per-layer adaptive injection: learned per-layer scaling for bigram residuals\n    3. Sinusoidal position enhancement: learned frequency scaling on sinusoidal \n       features mixed into learned position embeddings\n    \"\"\"\n    def __init__(self, config):\n        super().__init__()\n        self.wte = nn.Embedding(config.vocab_size, config.n_embd)\n        self.wpe = nn.Embedding(config.block_size, config.n_embd)\n        self.drop = nn.Dropout(config.dropout)\n        self.block_size = config.block_size\n        self.n_embd = config.n_embd\n        self.vocab_size = config.vocab_size\n        self.n_layer = config.n_layer\n        \n        # Bigram hash embedding: 5x vocab for hash collision reduction\n        self.bigram_vocab_size = config.vocab_size * 5\n        self.bigram_embed = nn.Embedding(self.bigram_vocab_size, config.n_embd)\n        nn.init.zeros_(self.bigram_embed.weight)\n        \n        # Per-layer scaling for bigram injection\n        self.bigram_lambdas = nn.Parameter(torch.full((config.n_layer,), 0.1))\n        \n        # Position-dependent gate for bigram injection (shape: block_size)\n        # Initialized to decay with position - more weight at early positions\n        pos_decay = torch.exp(-torch.arange(config.block_size, dtype=torch.float32) / (config.block_size / 4))\n        self.pos_gate = nn.Parameter(pos_decay)\n        \n        # Sinusoidal position enhancement: learnable frequency scaling\n        # Shape matches half the embedding dimension for proper sin/cos pairs\n        n_half = config.n_embd // 2\n        div_term = torch.exp(torch.arange(0, n_half).float() * -(math.log(10000.0) / n_half))\n        self.register_buffer('pos_div_term', div_term)\n        self.pos_freq_scale = nn.Parameter(torch.ones(n_half))\n        # Learnable mixing weight for sinusoidal enhancement (starts at 0)\n        self.sin_mix = nn.Parameter(torch.tensor(0.0))\n        \n        self._cached_bigram = None\n    \n    def _bigram_hash(self, idx):\n        \"\"\"Compute bigram hash indices from consecutive token pairs.\"\"\"\n        rand_int_1 = 36313\n        rand_int_2 = 27191\n        mod = self.bigram_vocab_size - 1\n        x = idx.to(torch.int32)\n        out = torch.zeros_like(x)\n        out[:, 0] = mod\n        out[:, 1:] = torch.bitwise_xor(\n            rand_int_1 * x[:, 1:],\n            rand_int_2 * x[:, :-1]\n        ) % mod\n        return out.long()\n    \n    def forward(self, idx):\n        b, t = idx.size()\n        tok_emb = self.wte(idx)\n        pos = torch.arange(0, t, dtype=torch.long, device=idx.device)\n        pos_emb = self.wpe(pos)\n        \n        # Compute and cache bigram embedding for per-layer injection\n        self._cached_bigram = self.bigram_embed(self._bigram_hash(idx))\n        \n        return self.drop(tok_emb + pos_emb)\n    \n    def get_value_embed(self, layer_idx):\n        \"\"\"Inject bigram embedding at every layer with position-aware scaling.\"\"\"\n        if self._cached_bigram is None or layer_idx >= self.n_layer:\n            return None\n        \n        # Position-dependent scaling: more weight at early positions\n        t = self._cached_bigram.size(1)\n        gate = self.pos_gate[:t]  # (T,)\n        scale = self.bigram_lambdas[layer_idx] * gate  # (T,)\n        \n        # Reshape to (1, T, 1) for broadcasting with (B, T, C)\n        return scale.view(1, t, 1) * self._cached_bigram\n    \n    def get_lm_head_weight(self):\n        return self.wte.weight\n    \n    def get_num_pos_params(self):\n        return self.wpe.weight.numel()"
7}