Upload train_bg.py with huggingface_hub

train_bg.py

@@ -0,0 +1,185 @@
+#!/usr/bin/env python3
+"""Background training for CogNet - resumes from checkpoint, trains N steps, saves."""
+import sys, os, time, math, random, json
+import torch
+import torch.nn as nn
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from cognet_1b import CogNet1B
+from infer import CharTokenizer
+
+CKPT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'checkpoints')
+N_STEPS = int(sys.argv[1]) if len(sys.argv) > 1 else 500
+LOG_FILE = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'train_bg.log')
+
+def log(msg):
+    with open(LOG_FILE, 'a') as f:
+        f.write(f'[{time.strftime("%H:%M:%S")}] {msg}\n')
+    print(msg, flush=True)
+
+# Generate training data from diverse text
+def make_training_data(tokenizer):
+    texts = []
+    # English texts
+    for _ in range(200):
+        texts.append("The quick brown fox jumps over the lazy dog. ")
+        texts.append("In the beginning there was the word and the word was good. ")
+        texts.append("Science tells us that the universe is vast and beautiful. ")
+        texts.append("Once upon a time there lived a king who ruled over a great kingdom. ")
+        texts.append("Knowledge is power and understanding is the key to wisdom. ")
+        texts.append("CogNet is a non-transformer language model with cognitive routing. ")
+        texts.append("The model processes sequences in linear time making it efficient. ")
+        texts.append("It has hierarchical memory with working episodic and semantic slots. ")
+        texts.append("This design enables contextual information without quadratic costs. ")
+        texts.append("The future of artificial intelligence is bright and full of possibilities. ")
+        texts.append("In the depth of winter I finally learned that within me there lay an invincible summer. ")
+        texts.append("The architecture consists of six adaptive computation blocks. ")
+        texts.append("Memory is organized in three tiers for efficient processing. ")
+        texts.append("CogNet replaces self-attention with O(n) cognitive routing. ")
+        texts.append("The model was trained on a CPU only machine demonstrating novel architectures. ")
+        texts.append("Hello world this is a test of the CogNet language model. ")
+        texts.append("The king sat on his throne and looked out over his kingdom with pride. ")
+        texts.append("A journey of a thousand miles begins with a single step forward. ")
+        texts.append("To be or not to be that is the question that puzzles many minds. ")
+        texts.append("The sun rose over the mountains casting long shadows across the valley. ")
+    # French texts
+    for _ in range(100):
+        texts.append("Bonjour le monde est beau et la science est merveilleuse. ")
+        texts.append("Le futur de l'intelligence artificielle est prometteur. ")
+        texts.append("Dans le profondeur de l'hiver j'ai appris qu'en moi habitait un été invincible. ")
+        texts.append("Le roi siégeait sur son trône et regardait son royaume avec fierté. ")
+        texts.append("La connaissance est le pouvoir et la compréhension est la clé de la sagesse. ")
+        texts.append("CogNet est un modèle de langage non-transformateur avec routage cognitif. ")
+        texts.append("Une voyage de mille lieues commence par un seul pas en avant. ")
+        texts.append("Le soleil s'est levé sur les montagnes projetant de longues ombres. ")
+        texts.append("La liberté est le droit de faire tout ce que les lois permettent. ")
+        texts.append("La musique est le langage universel de l'humanité toute entière. ")
+    
+    all_text = ''.join(texts)
+    ids = tokenizer.encode(all_text)
+    return torch.tensor(ids, dtype=torch.long)
+
+# Load tokenizer
+tokenizer = CharTokenizer.load(os.path.join(CKPT_DIR, 'tokenizer_v3.json'))
+log(f'Tokenizer: {tokenizer.vocab_size} chars')
+
+# Create model
+model = CogNet1B(vocab_size=tokenizer.vocab_size, hidden_dim=512, num_blocks=6,
+    num_channels=6, channel_dim=128, ff_dim=1024, routing_iters=1,
+    max_adaptive_steps=2, max_seq_len=192, working_slots=32,
+    episodic_slots=64, semantic_slots=128, key_dim=256, dropout=0.1)
+
+# Resume from checkpoint
+start_step = 0
+best_val = float('inf')
+ckpt_path = os.path.join(CKPT_DIR, 'cognet_best.pt')
+if os.path.exists(ckpt_path):
+    ckpt = torch.load(ckpt_path, map_location='cpu', weights_only=False)
+    state = ckpt['model_state_dict']
+    fp16_state = {k: v.float() if v.dtype == torch.float16 else v for k, v in state.items()}
+    model.load_state_dict(fp16_state)
+    start_step = ckpt.get('metrics', {}).get('step', 0)
+    best_val = ckpt.get('metrics', {}).get('val_loss', float('inf'))
+    log(f'Resumed from step {start_step}, best_val={best_val:.4f}')
+
+# Training data
+train_ids = make_training_data(tokenizer)
+log(f'Training data: {len(train_ids):,} tokens')
+
+# Optimizer
+optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=0.01, betas=(0.9, 0.95))
+COSINE_END = 10000
+
+def lr_lambda(step):
+    s = step + start_step
+    if s < 200: return s / 200
+    if s >= COSINE_END: return 0.05
+    return 0.5 * (1.0 + math.cos(math.pi * (s - 200) / max(COSINE_END - 200, 1)))
+
+scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+criterion = nn.CrossEntropyLoss(ignore_index=0)
+
+# Train
+model.train()
+SEQ_LEN, BS = 128, 2
+running_loss = 0.0
+max_start = len(train_ids) - SEQ_LEN - 1
+t0 = time.time()
+
+log(f'Starting training: {N_STEPS} steps from step {start_step}')
+
+for step in range(1, N_STEPS + 1):
+    optimizer.zero_grad(set_to_none=True)
+    acc = 0.0
+    for _ in range(4):  # gradient accumulation
+        starts = torch.randint(0, max_start, (BS,))
+        x = torch.stack([train_ids[s:s+SEQ_LEN] for s in starts])
+        y = torch.stack([train_ids[s+1:s+SEQ_LEN+1] for s in starts])
+        out = model(x)
+        loss = criterion(out['logits'].view(-1, out['logits'].size(-1)), y.view(-1))
+        (loss / 4).backward()
+        acc += loss.item()
+        del out, loss, x, y
+    
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+    optimizer.step()
+    scheduler.step()
+    running_loss += acc
+    gs = step + start_step
+    
+    if step % 25 == 0:
+        avg = running_loss / 25
+        ppl = math.exp(min(avg, 20))
+        lr = scheduler.get_last_lr()[0]
+        elapsed = time.time() - t0
+        spm = step / (elapsed / 60) if elapsed > 0 else 0
+        log(f'Step {gs:>5d} | Loss: {avg:.4f} | PPL: {ppl:.1f} | LR: {lr:.6f} | {spm:.1f} steps/min')
+        running_loss = 0.0
+    
+    # Save checkpoint every 200 steps
+    if step % 200 == 0:
+        # Quick eval
+        model.eval()
+        val_loss = 0.0
+        with torch.no_grad():
+            for _ in range(5):
+                s = torch.randint(0, max_start, (BS,))
+                x = torch.stack([train_ids[si:si+SEQ_LEN] for si in s])
+                y = torch.stack([train_ids[si+1:si+SEQ_LEN+1] for si in s])
+                out = model(x)
+                val_loss += criterion(out['logits'].view(-1, out['logits'].size(-1)), y.view(-1)).item()
+        val_loss /= 5
+        
+        is_best = val_loss < best_val
+        if is_best:
+            best_val = val_loss
+        
+        val_ppl = math.exp(min(val_loss, 20))
+        params = sum(p.numel() for p in model.parameters())
+        
+        torch.save({
+            'model_state_dict': model.state_dict(),
+            'vocab_size': tokenizer.vocab_size,
+            'hidden_dim': 512,
+            'num_blocks': 6,
+            'max_seq_len': 192,
+            'tokenizer_vocab_size': tokenizer.vocab_size,
+            'metrics': {'step': gs, 'val_loss': val_loss, 'val_ppl': val_ppl, 'total_params': params},
+        }, ckpt_path)
+        
+        log(f'Checkpoint saved: step={gs}, val_loss={val_loss:.4f}, val_ppl={val_ppl:.2f} {"NEW BEST!" if is_best else ""}')
+        model.train()
+
+# Final save
+gs = start_step + N_STEPS
+model.eval()
+params = sum(p.numel() for p in model.parameters())
+torch.save({
+    'model_state_dict': model.state_dict(),
+    'vocab_size': tokenizer.vocab_size,
+    'hidden_dim': 512,
+    'num_blocks': 6,
+    'max_seq_len': 192,
+    'tokenizer_vocab_size': tokenizer.vocab_size,
+    'metrics': {'step': gs, 'total_params': params},
+}, ckpt_path)
+log(f'Done: step {start_step}->{gs} in {time.time()-t0:.0f}s')