Upload 8 files

mixture_of_recursion.py

@@ -0,0 +1,418 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+from typing import Optional, Tuple
+import math
+
+@dataclass
+class RecursiveLanguageModelConfig:
+    vocab_size: int = 50257
+    embedding_dim: int = 512
+    num_layers: int = 6
+    num_attention_heads: int = 8
+    max_recursion_steps: int = 5
+    max_position_embeddings: int = 512
+    hidden_dropout_prob: float = 0.1
+    attention_dropout_prob: float = 0.1
+    intermediate_size: int = 2048
+    layer_norm_eps: float = 1e-5
+
+    pad_token_id: int = 50256
+    bos_token_id: int = 50256
+    eos_token_id: int = 50256
+
+    simple_recursion_steps: int = 1
+    medium_recursion_steps: int = 3
+    complex_recursion_steps: int = 5
+
+    confidence_threshold: float = 0.8
+    use_adaptive_stopping: bool = True
+    initializer_range: float = 0.02
+
+
+# Model Output class that supports subscripting
+class ModelOutput:
+    def __init__(self, loss=None, logits=None, complexity_class=None, recursion_steps=None):
+        self.loss = loss
+        self.logits = logits
+        self.complexity_class = complexity_class
+        self.recursion_steps = recursion_steps
+
+    def __getitem__(self, key):
+        if isinstance(key, str):
+            return getattr(self, key)
+        elif isinstance(key, int):
+            # For subscript access like outputs[0], outputs[1]
+            items = [self.loss, self.logits, self.complexity_class, self.recursion_steps]
+            return items[key]
+        elif isinstance(key, slice):
+            items = [self.loss, self.logits, self.complexity_class, self.recursion_steps]
+            return items[key]
+
+    def __iter__(self):
+        return iter([self.loss, self.logits, self.complexity_class, self.recursion_steps])
+
+
+class RotaryPositionalEmbedding(nn.Module):
+    def __init__(self, dim, max_seq_len=2048, base=10000):
+        super().__init__()
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer('inv_freq', inv_freq)
+        self.max_seq_len = max_seq_len
+        self.dim = dim
+
+    def forward(self, seq_len, device):
+        t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
+        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb.cos(), emb.sin()
+
+
+def apply_rotary_pos_emb(q, k, cos, sin):
+    def rotate_half(x):
+        x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
+        return torch.cat((-x2, x1), dim=-1)
+
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, config: RecursiveLanguageModelConfig):
+        super().__init__()
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.embedding_dim // config.num_attention_heads
+        self.embed_dim = config.embedding_dim
+
+        assert self.embed_dim % self.num_heads == 0
+
+        self.q_proj = nn.Linear(config.embedding_dim, config.embedding_dim)
+        self.k_proj = nn.Linear(config.embedding_dim, config.embedding_dim)
+        self.v_proj = nn.Linear(config.embedding_dim, config.embedding_dim)
+        self.out_proj = nn.Linear(config.embedding_dim, config.embedding_dim)
+
+        self.dropout = nn.Dropout(config.attention_dropout_prob)
+        self.rotary_emb = RotaryPositionalEmbedding(self.head_dim, config.max_position_embeddings)
+
+    def forward(self, hidden_states, attention_mask=None):
+        batch_size, seq_len, _ = hidden_states.shape
+
+        q = self.q_proj(hidden_states)
+        k = self.k_proj(hidden_states)
+        v = self.v_proj(hidden_states)
+
+        q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+        cos, sin = self.rotary_emb(seq_len, hidden_states.device)
+        cos = cos[None, None, :, :].expand(batch_size, self.num_heads, -1, -1)
+        sin = sin[None, None, :, :].expand(batch_size, self.num_heads, -1, -1)
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        if attention_mask is not None:
+            attn_weights = attn_weights + attention_mask
+
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+
+        attn_output = torch.matmul(attn_weights, v)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(batch_size, seq_len, self.embed_dim)
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output
+
+
+class FeedForward(nn.Module):
+    def __init__(self, config: RecursiveLanguageModelConfig):
+        super().__init__()
+        self.fc1 = nn.Linear(config.embedding_dim, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.embedding_dim)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.dropout(x)
+        x = self.fc2(x)
+        x = self.dropout(x)
+        return x
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: RecursiveLanguageModelConfig):
+        super().__init__()
+        self.attention = MultiHeadAttention(config)
+        self.feed_forward = FeedForward(config)
+        self.ln1 = nn.LayerNorm(config.embedding_dim, eps=config.layer_norm_eps)
+        self.ln2 = nn.LayerNorm(config.embedding_dim, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states, attention_mask=None):
+        residual = hidden_states
+        hidden_states = self.ln1(hidden_states)
+        hidden_states = self.attention(hidden_states, attention_mask)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.ln2(hidden_states)
+        hidden_states = self.feed_forward(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class SequenceLevelRouter(nn.Module):
+    def __init__(self, config: RecursiveLanguageModelConfig):
+        super().__init__()
+        self.config = config
+
+        self.pooler = nn.Linear(config.embedding_dim, config.embedding_dim)
+        self.pooler_activation = nn.Tanh()
+
+        self.classifier = nn.Sequential(
+            nn.Linear(config.embedding_dim, config.embedding_dim // 2),
+            nn.GELU(),
+            nn.Dropout(0.1),
+            nn.Linear(config.embedding_dim // 2, 3)
+        )
+
+    def forward(self, hidden_states, attention_mask=None):
+        if attention_mask is not None:
+            mask_expanded = attention_mask.unsqueeze(-1).float()
+            sum_hidden = torch.sum(hidden_states * mask_expanded, dim=1)
+            sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9)
+            pooled = sum_hidden / sum_mask
+        else:
+            pooled = hidden_states.mean(dim=1)
+
+        pooled = self.pooler(pooled)
+        pooled = self.pooler_activation(pooled)
+
+        complexity_logits = self.classifier(pooled)
+        complexity_class = torch.argmax(complexity_logits, dim=-1)
+
+        recursion_steps = torch.zeros_like(complexity_class)
+        recursion_steps[complexity_class == 0] = self.config.simple_recursion_steps
+        recursion_steps[complexity_class == 1] = self.config.medium_recursion_steps
+        recursion_steps[complexity_class == 2] = self.config.complex_recursion_steps
+
+        return complexity_logits, complexity_class, recursion_steps
+
+
+class RecursionLayer(nn.Module):
+    def __init__(self, config: RecursiveLanguageModelConfig):
+        super().__init__()
+        self.transformer_block = TransformerBlock(config)
+
+    def forward(self, hidden_states, attention_mask=None):
+        return self.transformer_block(hidden_states, attention_mask)
+
+
+class RecursiveLanguageModel(nn.Module):
+    def __init__(self, config: RecursiveLanguageModelConfig):
+        super().__init__()
+        self.config = config
+
+        self.embedding_layer = nn.Embedding(
+            config.vocab_size,
+            config.embedding_dim,
+            padding_idx=config.pad_token_id
+        )
+
+        self.base_transformer = nn.ModuleList([
+            TransformerBlock(config) for _ in range(config.num_layers)
+        ])
+
+        self.router = SequenceLevelRouter(config)
+        self.recursion_layer = RecursionLayer(config)
+
+        self.final_layer_norm = nn.LayerNorm(config.embedding_dim, eps=config.layer_norm_eps)
+        self.language_model_head = nn.Linear(config.embedding_dim, config.vocab_size, bias=False)
+
+        self.tie_weights()
+        self._init_weights()
+
+    def tie_weights(self):
+        self.language_model_head.weight = self.embedding_layer.weight
+
+    def _init_weights(self):
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+                if module.bias is not None:
+                    module.bias.data.zero_()
+            elif isinstance(module, nn.Embedding):
+                module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+                if module.padding_idx is not None:
+                    module.weight.data[module.padding_idx].zero_()
+            elif isinstance(module, nn.LayerNorm):
+                module.bias.data.zero_()
+                module.weight.data.fill_(1.0)
+
+    def get_attention_mask(self, input_ids):
+        batch_size, seq_len = input_ids.shape
+        device = input_ids.device
+
+        causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=device), diagonal=1).bool()
+        attention_mask = torch.zeros(batch_size, 1, seq_len, seq_len, device=device)
+        attention_mask[:, :, causal_mask] = float('-inf')
+
+        padding_mask = (input_ids == self.config.pad_token_id)
+        valid_mask = ~padding_mask
+
+        if padding_mask.any():
+            padding_mask_expanded = padding_mask.unsqueeze(1).unsqueeze(2)
+            attention_mask = attention_mask.masked_fill(padding_mask_expanded, float('-inf'))
+
+        return attention_mask, valid_mask
+
+    def forward(self, input_ids, labels=None, attention_mask=None):
+        batch_size, seq_len = input_ids.shape
+
+        hidden_states = self.embedding_layer(input_ids)
+        attn_mask, padding_mask = self.get_attention_mask(input_ids)
+
+        for layer in self.base_transformer:
+            hidden_states = layer(hidden_states, attn_mask)
+
+        complexity_logits, complexity_class, recursion_steps = self.router(
+            hidden_states, padding_mask
+        )
+
+        if self.training:
+            max_steps = self.config.complex_recursion_steps
+            for step in range(max_steps):
+                hidden_states = self.recursion_layer(hidden_states, attn_mask)
+        else:
+            max_steps_in_batch = int(recursion_steps.max().item())
+            for step in range(max_steps_in_batch):
+                step_mask = (recursion_steps > step).float().unsqueeze(-1).unsqueeze(-1)
+                new_hidden = self.recursion_layer(hidden_states, attn_mask)
+                hidden_states = step_mask * new_hidden + (1 - step_mask) * hidden_states
+
+        hidden_states = self.final_layer_norm(hidden_states)
+        logits = self.language_model_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            lm_loss = loss_fct(
+                shift_logits.view(-1, self.config.vocab_size),
+                shift_labels.view(-1)
+            )
+
+            complexity_value = min(max(seq_len // 170, 0), 2)
+            pseudo_labels = torch.full(
+                (batch_size,),
+                complexity_value,
+                dtype=torch.long,
+                device=input_ids.device
+            )
+
+            router_loss_fct = nn.CrossEntropyLoss()
+            router_loss = router_loss_fct(complexity_logits, pseudo_labels)
+
+            loss = lm_loss + 0.1 * router_loss
+
+        return ModelOutput(
+            loss=loss,
+            logits=logits,
+            complexity_class=complexity_class,
+            recursion_steps=recursion_steps
+        )
+
+    def generate(self, input_ids, max_new_tokens=50, temperature=1.0,
+                 top_p=0.9, do_sample=True):
+        self.eval()
+        generated = input_ids
+
+        for _ in range(max_new_tokens):
+            with torch.no_grad():
+                outputs = self.forward(generated)
+                logits = outputs.logits
+
+            next_token_logits = logits[:, -1, :] / temperature
+
+            if do_sample:
+                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                sorted_indices_to_remove[..., 0] = 0
+
+                indices_to_remove = sorted_indices_to_remove.scatter(
+                    1, sorted_indices, sorted_indices_to_remove
+                )
+                next_token_logits[indices_to_remove] = float('-inf')
+
+                probs = F.softmax(next_token_logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+
+            generated = torch.cat([generated, next_token], dim=-1)
+
+            if next_token.item() == self.config.eos_token_id:
+                break
+
+        return generated
+
+    def save_pretrained(self, save_directory):
+        import os
+        import json
+
+        os.makedirs(save_directory, exist_ok=True)
+        torch.save(self.state_dict(), os.path.join(save_directory, 'pytorch_model.bin'))
+
+        config_dict = {
+            'vocab_size': self.config.vocab_size,
+            'embedding_dim': self.config.embedding_dim,
+            'num_layers': self.config.num_layers,
+            'num_attention_heads': self.config.num_attention_heads,
+            'max_recursion_steps': self.config.max_recursion_steps,
+            'max_position_embeddings': self.config.max_position_embeddings,
+            'hidden_dropout_prob': self.config.hidden_dropout_prob,
+            'attention_dropout_prob': self.config.attention_dropout_prob,
+            'intermediate_size': self.config.intermediate_size,
+            'layer_norm_eps': self.config.layer_norm_eps,
+            'pad_token_id': self.config.pad_token_id,
+            'bos_token_id': self.config.bos_token_id,
+            'eos_token_id': self.config.eos_token_id,
+            'simple_recursion_steps': self.config.simple_recursion_steps,
+            'medium_recursion_steps': self.config.medium_recursion_steps,
+            'complex_recursion_steps': self.config.complex_recursion_steps,
+            'confidence_threshold': self.config.confidence_threshold,
+            'use_adaptive_stopping': self.config.use_adaptive_stopping,
+            'initializer_range': self.config.initializer_range,
+        }
+
+        with open(os.path.join(save_directory, 'config.json'), 'w') as f:
+            json.dump(config_dict, f, indent=2)
+
+    @classmethod
+    def from_pretrained(cls, load_directory, device='cpu'):
+        import os
+        import json
+
+        config_path = os.path.join(load_directory, 'config.json')
+        with open(config_path, 'r') as f:
+            config_dict = json.load(f)
+
+        config = RecursiveLanguageModelConfig(**config_dict)
+        model = cls(config)
+
+        weights_path = os.path.join(load_directory, 'pytorch_model.bin')
+        state_dict = torch.load(weights_path, map_location=device)
+        model.load_state_dict(state_dict)
+
+        model.to(device)
+        return model

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a940c0a5f02c276105651421d196092982e55bfb2b8d0c55240de58569a1a197
+size 192826915

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|endoftext|>",
+  "unk_token": "<|endoftext|>"
+}

tokenizer_config (2).json

@@ -0,0 +1,21 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "50256": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": true,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1024,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>"
+}

train (2).py

@@ -0,0 +1,252 @@
+import torch
+import torch.nn as nn
+import math
+from transformers import AutoTokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
+from datasets import load_dataset, interleave_datasets
+from mixture_of_recursion import RecursiveLanguageModel, RecursiveLanguageModelConfig
+import gc
+
+# Configuration
+TOTAL_SAMPLES = 50000
+BATCH_SIZE = 1
+GRAD_ACCUM = 32
+EPOCHS = 3
+LEARNING_RATE = 3e-4
+MAX_LENGTH = 384
+
+print("Starting training with 50K premium samples")
+print("-" * 60)
+
+# Load tokenizer
+print("\nLoading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+tokenizer.pad_token = tokenizer.eos_token
+
+print(f"Tokenizer vocab size: {len(tokenizer)}")
+print(f"Pad token ID: {tokenizer.pad_token_id}")
+
+# Load datasets
+print("\nLoading datasets...")
+print("  FineWeb-Edu (45%)")
+fineweb = load_dataset(
+    "HuggingFaceFW/fineweb-edu",
+    name="sample-10BT",
+    split="train",
+    streaming=True
+).shuffle(seed=42).take(int(TOTAL_SAMPLES * 0.45))
+
+print("  Cosmopedia (30%)")
+cosmopedia = load_dataset(
+    "HuggingFaceTB/cosmopedia",
+    "web_samples_v1",
+    split="train",
+    streaming=True
+).shuffle(seed=42).take(int(TOTAL_SAMPLES * 0.30))
+
+print("  OpenWebText (25%)")
+openwebtext = load_dataset(
+    "openwebtext",
+    split="train",
+    streaming=True
+).shuffle(seed=42).take(int(TOTAL_SAMPLES * 0.25))
+
+# Mix datasets
+print("\nMixing datasets...")
+train_dataset = interleave_datasets(
+    [fineweb, cosmopedia, openwebtext],
+    probabilities=[0.45, 0.30, 0.25],
+    seed=42
+)
+
+# Tokenization function
+def tokenize(examples):
+    if 'text' in examples:
+        texts = examples['text']
+    elif 'content' in examples:
+        texts = examples['content']
+    else:
+        texts = list(examples.values())[0]
+
+    return tokenizer(
+        texts,
+        truncation=True,
+        max_length=MAX_LENGTH,
+        padding=False
+    )
+
+# Tokenize datasets
+print("Tokenizing...")
+tokenized_train = train_dataset.map(
+    tokenize,
+    batched=True,
+    remove_columns=train_dataset.column_names
+).filter(lambda x: len(x['input_ids']) >= 128)
+
+# Validation set
+val_dataset = load_dataset(
+    "HuggingFaceFW/fineweb-edu",
+    name="sample-10BT",
+    split="train",
+    streaming=True
+).take(1000)
+
+val_tokenized = val_dataset.map(
+    tokenize,
+    batched=True,
+    remove_columns=val_dataset.column_names
+).filter(lambda x: len(x['input_ids']) >= 128)
+
+# Build model
+print("\nBuilding model...")
+config = RecursiveLanguageModelConfig(
+    vocab_size=len(tokenizer),
+    embedding_dim=512,
+    num_layers=6,
+    num_attention_heads=8,
+    max_recursion_steps=5,
+    max_position_embeddings=512,
+    intermediate_size=2048,
+    pad_token_id=tokenizer.pad_token_id,
+    bos_token_id=tokenizer.pad_token_id,
+    eos_token_id=tokenizer.pad_token_id,
+    simple_recursion_steps=1,
+    medium_recursion_steps=3,
+    complex_recursion_steps=5,
+    use_adaptive_stopping=True,
+    hidden_dropout_prob=0.1,
+    attention_dropout_prob=0.1
+)
+
+model = RecursiveLanguageModel(config)
+
+params = sum(p.numel() for p in model.parameters()) / 1e6
+print(f"Model parameters: {params:.1f}M")
+
+# Clear cache
+torch.cuda.empty_cache()
+gc.collect()
+
+# Training setup
+data_collator = DataCollatorForLanguageModeling(
+    tokenizer=tokenizer,
+    mlm=False
+)
+
+steps_per_epoch = TOTAL_SAMPLES // (BATCH_SIZE * GRAD_ACCUM)
+max_steps = steps_per_epoch * EPOCHS
+
+print(f"\nTraining steps: {max_steps}")
+print(f"Effective batch size: {BATCH_SIZE * GRAD_ACCUM}")
+
+training_args = TrainingArguments(
+    output_dir="./checkpoints",
+    max_steps=max_steps,
+    per_device_train_batch_size=BATCH_SIZE,
+    per_device_eval_batch_size=BATCH_SIZE,
+    gradient_accumulation_steps=GRAD_ACCUM,
+    learning_rate=LEARNING_RATE,
+    weight_decay=0.01,
+    warmup_steps=500,
+    fp16=True,
+    logging_steps=100,
+    eval_strategy="steps",
+    eval_steps=1000,
+    save_steps=1000,
+    save_total_limit=2,
+    load_best_model_at_end=True,
+    metric_for_best_model="eval_loss",
+    report_to="none",
+    max_grad_norm=1.0,
+    save_safetensors=False,  # Use PyTorch format instead of safetensors
+)
+
+# Custom trainer with perplexity
+class CustomTrainer(Trainer):
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        outputs = model(**inputs)
+        return (outputs.loss, outputs) if return_outputs else outputs.loss
+
+    def evaluation_loop(self, dataloader, description, prediction_loss_only=None, 
+                       ignore_keys=None, metric_key_prefix="eval"):
+        output = super().evaluation_loop(
+            dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix
+        )
+        
+        if output.metrics.get(f"{metric_key_prefix}_loss") is not None:
+            try:
+                perplexity = math.exp(output.metrics[f"{metric_key_prefix}_loss"])
+                output.metrics[f"{metric_key_prefix}_perplexity"] = perplexity
+            except OverflowError:
+                output.metrics[f"{metric_key_prefix}_perplexity"] = float("inf")
+        
+        return output
+    
+    def training_step(self, model, inputs, num_items_in_batch=None):
+        loss = super().training_step(model, inputs, num_items_in_batch)
+        
+        if self.state.global_step % 50 == 0:
+            torch.cuda.empty_cache()
+        
+        return loss
+
+trainer = CustomTrainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_train,
+    eval_dataset=val_tokenized,
+    data_collator=data_collator
+)
+
+# Train
+print("\nStarting training...")
+print("-" * 60)
+
+try:
+    trainer.train()
+    
+    # Final evaluation
+    print("\nFinal evaluation...")
+    metrics = trainer.evaluate()
+    
+    print("\n" + "="*60)
+    print("FINAL RESULTS:")
+    print("="*60)
+    print(f"Evaluation Loss: {metrics['eval_loss']:.4f}")
+    
+    if 'eval_perplexity' in metrics:
+        print(f"Perplexity: {metrics['eval_perplexity']:.2f}")
+    else:
+        try:
+            perplexity = math.exp(metrics['eval_loss'])
+            print(f"Perplexity: {perplexity:.2f}")
+        except OverflowError:
+            print(f"Perplexity: inf (loss too high)")
+    print("="*60 + "\n")
+    
+    # Save with custom method (handles tied weights properly)
+    print("Saving model...")
+    model.save_pretrained("./recursive-lm")
+    tokenizer.save_pretrained("./recursive-lm")
+    print("Model saved successfully!")
+    
+except KeyboardInterrupt:
+    print("\n\nTraining interrupted by user")
+    print("Saving current model state...")
+    model.save_pretrained("./recursive-lm-interrupted")
+    tokenizer.save_pretrained("./recursive-lm-interrupted")
+    
+except Exception as e:
+    print(f"\n\nTraining stopped due to: {e}")
+    import traceback
+    traceback.print_exc()
+    
+    # Try to save anyway
+    try:
+        print("\nAttempting to save model...")
+        model.save_pretrained("./recursive-lm-error")
+        tokenizer.save_pretrained("./recursive-lm-error")
+        print("Model saved!")
+    except:
+        print("Could not save model")
+
+print("\nTraining complete!")