!pip install safetensors -q

import time
import torch
import torch.nn as nn
from transformers import AutoTokenizer
from datasets import load_dataset
from torch.utils.data import DataLoader
from safetensors.torch import save_file

print("🚀 Initializing TINY Model Dual-T4 Master Script...")

# 1. Dataset & Tokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token 

print("Downloading Clean TinyStories Dataset...")
dataset = load_dataset("roneneldan/TinyStories", split="train")
small_dataset = dataset.select(range(100000))

def tokenize_function(example):
    tokens = tokenizer(example['text'], truncation=True, max_length=512, padding="max_length", return_tensors="pt")
    return {"input_ids": tokens["input_ids"][0]}

print("Tokenizing Dataset...")
tokenized_dataset = small_dataset.map(tokenize_function, remove_columns=dataset.column_names, num_proc=4)
tokenized_dataset.set_format("torch")
train_dataloader = DataLoader(tokenized_dataset, batch_size=16, shuffle=True)

# 2. Architecture (5M Params)
class TinyLLM(nn.Module):
    def __init__(self, vocab_size=50257, d_model=128, n_heads=4, n_layers=4, max_seq_len=512):
        super().__init__()
        self.d_model = d_model
        self.token_emb = nn.Embedding(vocab_size, d_model)
        self.pos_emb = nn.Embedding(max_seq_len, d_model)
        
        decoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=n_heads, dim_feedforward=4 * d_model, 
            batch_first=True, activation="gelu"
        )
        self.transformer = nn.TransformerEncoder(decoder_layer, num_layers=n_layers)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        self.token_emb.weight = self.lm_head.weight

    # Loss calculation inside forward to prevent GPU 0 OOM
    def forward(self, x, labels=None):
        seq_len = x.size(1)
        mask = nn.Transformer.generate_square_subsequent_mask(seq_len).to(x.device)
        positions = torch.arange(0, seq_len, dtype=torch.long, device=x.device).unsqueeze(0)
        x = self.token_emb(x) + self.pos_emb(positions)
        x = self.transformer(x, mask=mask, is_causal=True)
        logits = self.lm_head(x)
        
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
            return loss
        return logits

# 3. Training Setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TinyLLM()

if torch.cuda.device_count() > 1:
    print(f"🔥 Detected {torch.cuda.device_count()} GPUs! Wrapping model...")
    model = nn.DataParallel(model)

model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4)
scaler = torch.amp.GradScaler('cuda')

model.train()
print("🔥 Starting 50M Token Training on 2x T4...")

start_time = time.time()
for step, batch in enumerate(train_dataloader):
    inputs = batch['input_ids'].to(device)
    
    optimizer.zero_grad()
    
    with torch.amp.autocast('cuda'):
        loss = model(inputs, labels=inputs).mean()
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    
    if step % 50 == 0 and step > 0:
        elapsed = time.time() - start_time
        steps_per_sec = 50 / elapsed
        print(f"Step {step} | Loss: {loss.item():.4f} | Speed: {steps_per_sec:.2f} steps/sec")
        start_time = time.time()

print("✅ Dual T4 Training Complete! Saving model...")

# 4. Save Block (Both Formats)
# Remove DataParallel wrapper for saving pure weights
model_to_save = model.module if hasattr(model, 'module') else model

# PyTorch Default Format
torch.save(model_to_save.state_dict(), 'tiny_model_t4.pth')
print("✅ Saved PyTorch format: 'tiny_model_t4.pth'")

# Safetensors Format
save_file(model_to_save.state_dict(), 'tiny_model_t4.safetensors')
print("✅ Saved Safetensors format: 'tiny_model_t4.safetensors'")
print("🎉 All Done! Download your files from the right panel.")