Upload transformerdecoder.py

transformerdecoder.py

@@ -0,0 +1,248 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+import numpy as np
+import math
+from transformers import AutoTokenizer, PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions
+
+import torchvision
+from torch.utils.data import Dataset, DataLoader
+from datasets import load_dataset_builder
+from datasets import load_dataset
+from transformers import DataCollatorForLanguageModeling
+from transformers import DataCollatorWithPadding, Trainer, TrainingArguments
+from torch.optim import AdamW
+from trl import SFTTrainer, SFTConfig
+from transformers import TrainingArguments, Trainer
+
+pretrain_data = load_dataset("Salesforce/wikitext", "wikitext-103-v1", split="train") # ["text"] contains the data
+tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")
+vocab_size = len(tokenizer)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token   
+
+
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model):
+        super().__init__()
+        self.pos_enc = nn.Sequential(
+            nn.Linear(1, d_model*4),
+            nn.Tanh(),
+            nn.Linear(d_model*4, d_model)
+        )
+    def forward(self, seq_len, device):
+        pos = torch.arange(seq_len, device=device, dtype=torch.float32).unsqueeze(-1) # (seq_len, 1)
+        pe = self.pos_enc(pos) # (seq_len, d_model)
+        return pe.unsqueeze(0) # (1, seq_len, d_model)
+    
+class AugmentedPositionGPTConfig(PretrainedConfig):
+    model_type = "AugmentedPositionGPT"
+
+    def __init__(
+        self,
+        vocab_size=vocab_size,
+        d_model=128,
+        num_heads=2,
+        num_layers=1,
+        max_position_embeddings=512,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.max_position_embeddings = max_position_embeddings
+
+
+
+
+
+
+class AugmentedPositionGPTBlock(nn.Module):
+    def __init__(self, d_model, num_heads):
+        super().__init__()
+        self.d_model = d_model
+        #self.output_embedding = nn.Embedding(vocab_size, d_model)
+        self.multiheadattention = nn.MultiheadAttention(d_model, num_heads, batch_first=True)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.ffn1 = nn.Linear(d_model, 4*d_model)
+        self.ffn2 = nn.Linear(d_model*4, d_model)
+        #self.linear = nn.Linear(d_model, vocab_size)
+    def forward(self, x, causal_mask=None):
+        residual = x
+        normx = self.norm1(x)
+        attn_out, _ = self.multiheadattention(normx, normx, normx, attn_mask=causal_mask) # Attention(Q, K, V) = softmax(Q @ K.T / sqrt(d_k) + mask) @ V
+        # output: (batch, seq_len, d_model)
+        x = residual + attn_out
+        residual2 = x
+        j = self.ffn1(self.norm2(x)) # takes in: (batch, seq_len, d_model)
+        
+        
+        # outputs: (batch, seq_len, d_model*4)
+        h = self.ffn2(F.relu(j)) # takes in: (batch, seq_len, d_model*4)
+        x = residual2 + h
+        # outputs: (batch, seq_len, d_model)
+
+        return x
+class AugmentedPositionGPT(PreTrainedModel):
+    config_class = AugmentedPositionGPTConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.vocab_size = config.vocab_size
+        self.d_model = config.d_model
+        self.num_heads = config.num_heads
+        self.num_layers = config.num_layers
+        self.max_position_embeddings = config.max_position_embeddings
+        self.output_embedding = nn.Embedding(self.vocab_size, self.d_model)
+        self.blocks = nn.ModuleList(
+            [AugmentedPositionGPTBlock(self.d_model, self.num_heads) for _ in range(self.num_layers)]
+        )
+        self.ln_f = nn.LayerNorm(self.d_model)
+        self.register_buffer(
+            "position_ids",
+            torch.arange(self.max_position_embeddings).unsqueeze(0), # (1, seq_len)
+            persistent=False,
+        )
+        self.post_init()
+
+    def causal_mask(self, seq_len, device):
+        # (seq_len, seq_len)
+        causal_mask = nn.Transformer.generate_square_subsequent_mask(seq_len, device=device)
+        return causal_mask
+    
+    def positional_encoding(self, seq_len, device):
+        d_model = self.d_model
+        # EVEN: PE(pos, 2i) = sin(pos/10000^(2i/d_model))
+        # ODD:  PE(pos, 2i+1) = cos(pos/10000^(2i/dmodel))
+        a = 10000
+
+        
+        i = torch.arange(0, d_model, 2, device=device, dtype=torch.float32) # (d_model/2)
+        div_term = a ** (i / d_model) # (d_model/2)
+        position = torch.arange(seq_len, device=device, dtype=torch.float32).unsqueeze(1) # (seq_len, 1)
+        angles = position / div_term # (seq_len, d_model/2)
+        pe = torch.zeros(seq_len, d_model, device=device, dtype=torch.float32) # (seq_len, d_model)
+        pe[:, 0::2] = torch.sin(angles)
+        pe[:, 1::2] = torch.cos(angles)
+        pe = pe.unsqueeze(0)
+        # shape: (1, seq_len, d_model)
+        return pe
+
+    def forward(
+            self,
+            input_ids=None,
+            attention_mask=None,
+            input_embeds=None,
+            output_hidden_states=False,
+            return_dict=True
+    ):
+        if input_ids is not None and input_embeds is not None:
+            raise ValueError("you cant specify both input_ids and input_embeds")
+        if input_embeds is None:
+            #max_id = input_ids.max().item()
+            #min_id = input_ids.min().item()
+            #if max_id >= self.vocab_size or min_id < 0:
+                #raise RuntimeError(
+                    #f"Bad token id: min={min_id}, max={max_id}, "
+                    #f"embedding vocab_size={self.vocab_size}"
+                #)
+            input_embeds = self.output_embedding(input_ids) # (batch, seq_len, d_model)
+        batch, seq_len, _ = input_embeds.shape
+        device = input_embeds.device
+        # output embeddings and postional encoding
+        x = self.output_embedding(input_ids) # (batch, seq_len, d_model)
+        pe = self.positional_encoding(seq_len, device=device) # (1, seq_len, d_model)
+        x = x + pe # (batch, seq_len, d_model)
+        causal_mask = self.causal_mask(seq_len, device)
+        all_hidden_states = [] if output_hidden_states else None
+        for block in self.blocks:
+            if output_hidden_states:
+                all_hidden_states.append(x)
+            x = block(x, causal_mask=causal_mask)
+        x = self.ln_f(x)
+        if not return_dict:
+            return (x, all_hidden_states)
+        
+        return {"last_hidden_state": x, "hidden_states": all_hidden_states}
+    
+
+class AugmentedPositionGPTForCausalLM(PreTrainedModel):
+    config_class = AugmentedPositionGPTConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformerdecoder = AugmentedPositionGPT(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+
+        self.lm_head.weight = self.transformerdecoder.output_embedding.weight
+        self._dynamic_tied_weights_keys = { # make sure to tell huggingface everything you do or else it will explode
+            "lm_head.weight": "transformerdecoder.output_embedding.weight"
+        }
+        self.post_init()
+    
+    
+
+    def forward(self, input_ids=None, attention_mask=None, input_embeds=None, labels=None, output_hidden_states=False, return_dict=True):
+        outputs= self.transformerdecoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            input_embeds = input_embeds,
+            output_hidden_states=output_hidden_states,
+            return_dict=True
+        )
+        hidden_states = outputs["last_hidden_state"] # (batch, seq_len, d_model)
+        logits = self.lm_head(hidden_states) # (batch, seq_len, vocab_size)
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                logits.view(-1, logits.size(-1)),
+                labels.view(-1)
+            )
+        if not return_dict:
+            output = (logits,)
+            if output_hidden_states:
+                output += (outputs["hidden_states"],)
+            return ((loss,) + output) if loss is not None else output
+        return CausalLMOutputWithCrossAttentions(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs["hidden_states"],
+            attentions=None,
+            cross_attentions=None
+        )
+config = AugmentedPositionGPTConfig(vocab_size=vocab_size)
+collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+model = AugmentedPositionGPTForCausalLM(config)
+def tokenize(examples):
+    return tokenizer(examples["text"], truncation=True, max_length=512)
+pretrain_data_tok = pretrain_data.map(
+    tokenize,
+    batched=True,
+    remove_columns=["text"],  # remove raw text so Trainer doesn't pass it
+)
+
+training_args = TrainingArguments(
+    output_dir = "AugmentedGPT/results",
+    num_train_epochs=1,
+    per_device_eval_batch_size=1,
+    remove_unused_columns=False,
+    gradient_accumulation_steps=8,
+    fp16=True,
+    
+
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=pretrain_data_tok,
+    data_collator=collator,
+)
+trainer.train()