Upload 2 files

image_model.py

@@ -0,0 +1,365 @@
+"""
+image_model.py — A real CNN (Convolutional Neural Network) built from scratch in PyTorch.
+
+Architecture:
+  Input (3 x 64 x 64 image)
+      → Conv2d(3, 32, 3)  + BatchNorm + ReLU + MaxPool    →  32 x 31 x 31
+      → Conv2d(32, 64, 3) + BatchNorm + ReLU + MaxPool    →  64 x 14 x 14
+      → Conv2d(64,128, 3) + BatchNorm + ReLU + MaxPool    → 128 x  6 x  6
+      → Flatten → Linear(128*6*6, 512) → ReLU → Dropout
+      → Linear(512, 128) → ReLU
+      → Linear(128, 8 categories)
+
+This CNN learns to LOOK at images the same way your text network learns to READ text.
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import threading
+_CNN_LOCK = threading.Lock()
+import numpy as np
+import os
+import json
+from datetime import datetime, UTC
+from pathlib import Path
+from collections import defaultdict
+
+try:
+    from PIL import Image
+    PIL_AVAILABLE = True
+except ImportError:
+    PIL_AVAILABLE = False
+
+IMG_SIZE   = 64       # resize all images to 64x64 (small = faster on CPU)
+CATEGORIES = [
+    'nature', 'technology', 'science', 'people',
+    'animals', 'food', 'sports', 'architecture'
+]
+
+CNN_CHECKPOINT  = 'cnn_checkpoint.pt'
+CNN_STATS_FILE  = 'cnn_stats.json'
+
+
+# ── IMAGE PREPROCESSING ───────────────────────────────────────────────────────
+def load_image_tensor(path: str, size: int = IMG_SIZE):
+    """Load image from disk → normalised float tensor (3, size, size)."""
+    if not PIL_AVAILABLE:
+        raise RuntimeError("Pillow not installed. Add 'Pillow' to requirements.txt")
+    img = Image.open(path).convert('RGB')
+    img = img.resize((size, size), Image.BILINEAR)
+    arr = np.array(img, dtype=np.float32) / 255.0
+    # Normalize with ImageNet mean/std (works well even for non-ImageNet data)
+    mean = np.array([0.485, 0.456, 0.406])
+    std  = np.array([0.229, 0.224, 0.225])
+    arr  = (arr - mean) / std
+    return torch.tensor(arr).permute(2, 0, 1)   # HWC → CHW
+
+
+# ── CNN MODEL ─────────────────────────────────────────────────────────────────
+class ImageCNN(nn.Module):
+    """
+    A real Convolutional Neural Network.
+    Learns to detect edges → shapes → textures → objects, layer by layer.
+    Each Conv2d layer is looking for patterns the previous layer found.
+    """
+
+    def __init__(self, num_classes: int = 8):
+        super().__init__()
+        self.num_classes = num_classes
+
+        # Convolutional feature extractor
+        self.features = nn.Sequential(
+            # Block 1 — learns basic edges and colours
+            nn.Conv2d(3, 32, kernel_size=3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+            nn.Conv2d(32, 32, kernel_size=3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),        # 64→32
+            nn.Dropout2d(0.1),
+
+            # Block 2 — learns corners, curves, textures
+            nn.Conv2d(32, 64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.Conv2d(64, 64, kernel_size=3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),        # 32→16
+            nn.Dropout2d(0.15),
+
+            # Block 3 — learns complex shapes and object parts
+            nn.Conv2d(64, 128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.Conv2d(128, 128, kernel_size=3, padding=1),
+            nn.BatchNorm2d(128),
+            nn.ReLU(),
+            nn.MaxPool2d(2, 2),        # 16→8
+            nn.Dropout2d(0.2),
+        )
+
+        # Classifier head
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(128 * 8 * 8, 512),
+            nn.ReLU(),
+            nn.Dropout(0.4),
+            nn.Linear(512, 128),
+            nn.ReLU(),
+            nn.Linear(128, num_classes),
+        )
+
+        # Activation tracking for visualization
+        self._activations = {}
+        self._register_hooks()
+        self._initialize_weights()
+
+    def _initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                nn.init.constant_(m.bias, 0)
+
+    def _register_hooks(self):
+        def make_hook(name):
+            def hook(module, inp, out):
+                if isinstance(out, torch.Tensor):
+                    v = out.detach().float()
+                    if v.dim() > 1:
+                        v = v.mean(0)
+                    if v.dim() > 1:
+                        v = v.mean(-1).mean(-1)  # spatial mean for conv layers
+                    self._activations[name] = v[:16].tolist()
+            return hook
+        for i, layer in enumerate(self.features):
+            layer.register_forward_hook(make_hook(f'conv_{i}'))
+        for i, layer in enumerate(self.classifier):
+            layer.register_forward_hook(make_hook(f'fc_{i}'))
+
+    def forward(self, x):
+        x = self.features(x)
+        return self.classifier(x)
+
+    def get_activations(self):
+        return dict(self._activations)
+
+    def get_feature_maps(self, x):
+        """Return intermediate feature maps for visualization."""
+        maps = {}
+        for i, layer in enumerate(self.features):
+            x = layer(x)
+            if isinstance(layer, nn.ReLU):
+                maps[f'relu_{i}'] = x.detach()
+        return maps
+
+
+# ── LIVING IMAGE NETWORK ──────────────────────────────────────────────────────
+class LivingImageNetwork:
+    """
+    Wraps the CNN with training loop, data loading, and stats.
+    Trains on images downloaded by ImageFetcher.
+    """
+
+    def __init__(self):
+        self.model      = ImageCNN(num_classes=len(CATEGORIES))
+        self.optimizer  = optim.Adam(self.model.parameters(), lr=0.001, weight_decay=1e-4)
+        self.scheduler  = optim.lr_scheduler.StepLR(self.optimizer, step_size=50, gamma=0.8)
+        self.criterion  = nn.CrossEntropyLoss()
+
+        self.epoch           = 0
+        self.total_images    = 0
+        self.loss_history    = []
+        self.acc_history     = []
+        self.category_counts = defaultdict(int)
+
+        self.stats = {
+            'epoch':        0,
+            'loss':         '—',
+            'accuracy':     '—',
+            'total_images': 0,
+            'lr':           0.001,
+            'last_image':   '(none yet)',
+            'status':       'idle',
+        }
+
+        self._load_checkpoint()
+
+    # ── DATA LOADING ──────────────────────────────────────────────────────────
+    def _load_batch(self, image_dir: Path, batch_size: int = 16):
+        """
+        Load a random batch of images from image_data/ folder.
+        Returns (tensor_batch, label_batch) or None if not enough images.
+        """
+        if not PIL_AVAILABLE:
+            return None
+
+        all_paths = []
+        for cat_idx, cat in enumerate(CATEGORIES):
+            cat_dir = image_dir / cat
+            if cat_dir.exists():
+                for p in cat_dir.iterdir():
+                    if p.suffix.lower() in ('.jpg', '.jpeg', '.png', '.webp'):
+                        all_paths.append((str(p), cat_idx))
+
+        if len(all_paths) < batch_size:
+            return None
+
+        import random
+        batch_paths = random.sample(all_paths, batch_size)
+        tensors, labels = [], []
+
+        for path, label in batch_paths:
+            try:
+                t = load_image_tensor(path)
+                tensors.append(t)
+                labels.append(label)
+                self.category_counts[CATEGORIES[label]] += 1
+            except Exception:
+                continue
+
+        if not tensors:
+            return None
+
+        return torch.stack(tensors), torch.tensor(labels, dtype=torch.long)
+
+    # ── TRAINING ──────────────────────────────────────────────────────────────
+    def train_step(self, image_dir: Path, batch_size: int = 16):
+        """One training step — load images, forward pass, backprop."""
+        batch = self._load_batch(image_dir, batch_size)
+        if batch is None:
+            return None
+
+        x, y = batch
+        x = x.float()  # ensure float32
+        with _CNN_LOCK:
+            self.model.train()
+            self.model.zero_grad(set_to_none=True)
+            logits = self.model(x)
+            loss   = self.criterion(logits, y)
+            loss.backward()
+            for p in self.model.parameters():
+                if p.grad is not None:
+                    p.grad.data.clamp_(-1.0, 1.0)
+            self.optimizer.step()
+            loss_val = loss.detach().item()
+            acc      = (logits.detach().argmax(1) == y).float().mean().item()
+
+        self.epoch       += 1
+        self.total_images += len(x)
+        self.scheduler.step()
+
+        self.loss_history.append(round(loss_val, 5))
+        self.acc_history.append(round(acc, 4))
+        if len(self.loss_history) > 500:
+            self.loss_history = self.loss_history[-500:]
+            self.acc_history  = self.acc_history[-500:]
+
+        last_path = batch[0]  # just the paths string
+        self.stats.update({
+            'epoch':        self.epoch,
+            'loss':         round(loss_val, 4),
+            'accuracy':     round(acc * 100, 1),
+            'total_images': self.total_images,
+            'lr':           round(self.optimizer.param_groups[0]['lr'], 7),
+        })
+
+        if self.epoch % 20 == 0:
+            self._save_checkpoint()
+        self._write_stats()
+        return loss_val
+
+    def train_n_steps(self, image_dir: Path, n: int = 20):
+        losses = []
+        for _ in range(n):
+            l = self.train_step(image_dir)
+            if l is not None:
+                losses.append(l)
+        return {
+            'steps':    len(losses),
+            'avg_loss': round(sum(losses)/len(losses), 5) if losses else None,
+        }
+
+    # ── INFERENCE ─────────────────────────────────────────────────────────────
+    def predict_image(self, image_path: str) -> dict:
+        """Predict category of a single image."""
+        if not PIL_AVAILABLE:
+            return {'error': 'Pillow not installed'}
+        try:
+            t = load_image_tensor(image_path).unsqueeze(0)
+            self.model.eval()
+            with torch.no_grad():
+                logits = self.model(t)
+                probs  = torch.softmax(logits, dim=1)[0].tolist()
+                pred   = int(logits.argmax(1).item())
+            return {
+                'prediction': CATEGORIES[pred],
+                'confidence': round(probs[pred] * 100, 1),
+                'all_probs':  {c: round(p*100, 2) for c, p in zip(CATEGORIES, probs)},
+            }
+        except Exception as e:
+            return {'error': str(e)}
+
+    # ── VIZ STATE ─────────────────────────────────────────────────────────────
+    def get_viz_state(self) -> dict:
+        self.model.eval()
+        with torch.no_grad():
+            dummy = torch.zeros(1, 3, IMG_SIZE, IMG_SIZE)
+            self.model(dummy)
+        return {
+            'layer_sizes':  [3, 32, 64, 128, 512, 128, len(CATEGORIES)],
+            'activations':  self.model.get_activations(),
+            'loss_history': self.loss_history[-100:],
+            'acc_history':  self.acc_history[-100:],
+            'stats':        self.stats,
+            'type':         'cnn',
+        }
+
+    # ── PERSISTENCE ───────────────────────────────────────────────────────────
+    def _save_checkpoint(self):
+        try:
+            torch.save({
+                'model':           self.model.state_dict(),
+                'optimizer':       self.optimizer.state_dict(),
+                'epoch':           self.epoch,
+                'total_images':    self.total_images,
+                'loss_history':    self.loss_history,
+                'acc_history':     self.acc_history,
+                'category_counts': dict(self.category_counts),
+            }, CNN_CHECKPOINT)
+        except Exception:
+            pass
+
+    def _load_checkpoint(self):
+        if not os.path.exists(CNN_CHECKPOINT):
+            return
+        try:
+            ck = torch.load(CNN_CHECKPOINT, map_location='cpu')
+            self.model.load_state_dict(ck['model'])
+            self.optimizer.load_state_dict(ck['optimizer'])
+            self.epoch           = ck.get('epoch', 0)
+            self.total_images    = ck.get('total_images', 0)
+            self.loss_history    = ck.get('loss_history', [])
+            self.acc_history     = ck.get('acc_history', [])
+            self.category_counts = defaultdict(int, ck.get('category_counts', {}))
+            if self.epoch > 0:
+                self.stats.update({
+                    'epoch':        self.epoch,
+                    'total_images': self.total_images,
+                })
+        except Exception:
+            pass
+
+    def _write_stats(self):
+        try:
+            with open(CNN_STATS_FILE, 'w') as f:
+                json.dump(self.stats, f, indent=2)
+        except Exception:
+            pass

neural_network.py

@@ -0,0 +1,441 @@
+"""
+neural_network.py — Real PyTorch neural network built from scratch.
+
+KEY CHANGES:
+  - Every item ingested is IMMEDIATELY written to knowledge.jsonl
+  - On startup, knowledge.jsonl is read back → data_buffer is restored
+  - Model checkpoint auto-saves every 30 training epochs
+  - training_stats.json written after every training step (human-readable)
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import threading
+_TEXT_LOCK = threading.Lock()
+import numpy as np
+import os
+import json
+import re
+from datetime import datetime, UTC
+from collections import Counter, defaultdict
+
+# ─── FILES ON DISK ────────────────────────────────────────────────────────────
+KNOWLEDGE_FILE   = 'knowledge.jsonl'      # Every article/text the AI has seen
+CHECKPOINT_FILE  = 'model_checkpoint.pt'  # PyTorch weights + optimizer state
+STATS_FILE       = 'training_stats.json'  # Human-readable live stats
+
+# ─── CATEGORIES ───────────────────────────────────────────────────────────────
+CATEGORIES = ['technology', 'science', 'world', 'sports',
+              'business', 'health', 'entertainment', 'other']
+
+# ─── VOCABULARY ───────────────────────────────────────────────────────────────
+STOPWORDS = {
+    'a','an','the','is','it','in','on','at','to','for','of','and','or','but',
+    'was','are','were','be','been','have','has','had','do','does','did','will',
+    'would','could','should','may','might','that','this','these','those','with',
+    'from','by','as','not','also','than','then','so','if','when','what','how',
+    'who','which','its','their','our','your','my','his','her','we','they','he',
+    'she','you','i','me','him','us','them','said','says','new','one','two',
+}
+
+class Vocabulary:
+    def __init__(self, max_size=10000):
+        self.word2idx   = {'<PAD>': 0, '<UNK>': 1}
+        self.idx2word   = {0: '<PAD>', 1: '<UNK>'}
+        self.word_counts = Counter()
+        self.max_size   = max_size
+        self.is_built   = False
+
+    def update(self, text: str):
+        self.word_counts.update(self._tokenize(text))
+
+    def build(self):
+        top = self.word_counts.most_common(self.max_size - 2)
+        self.word2idx = {'<PAD>': 0, '<UNK>': 1}
+        self.idx2word = {0: '<PAD>', 1: '<UNK>'}
+        for i, (word, _) in enumerate(top):
+            idx = i + 2
+            self.word2idx[word] = idx
+            self.idx2word[idx]  = word
+        self.is_built = True
+
+    def encode(self, text: str, max_len: int = 64) -> list:
+        words = self._tokenize(text)[:max_len]
+        ids   = [self.word2idx.get(w, 1) for w in words]
+        ids  += [0] * (max_len - len(ids))
+        return ids
+
+    def _tokenize(self, text: str) -> list:
+        text = text.lower()
+        text = re.sub(r'[^\w\s]', ' ', text)
+        return [w for w in text.split() if w not in STOPWORDS and len(w) > 2]
+
+    def __len__(self):
+        return len(self.word2idx)
+
+
+# ─── MODEL ────────────────────────────────────────────────────────────────────
+class TextClassifier(nn.Module):
+    def __init__(self, vocab_size=10002, embed_dim=64,
+                 hidden=[256, 128, 64], num_classes=8):
+        super().__init__()
+        self.embed_dim   = embed_dim
+        self.hidden_dims = hidden
+        self.num_classes = num_classes
+
+        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
+        nn.init.normal_(self.embedding.weight, 0, 0.1)
+
+        layers = []
+        in_dim = embed_dim
+        for h in hidden:
+            layers += [nn.Linear(in_dim, h), nn.LayerNorm(h),
+                       nn.ReLU(), nn.Dropout(0.25)]
+            in_dim = h
+        layers.append(nn.Linear(in_dim, num_classes))
+        self.net = nn.Sequential(*layers)
+
+        self._activations = {}
+        self._register_hooks()
+
+    def _register_hooks(self):
+        def make_hook(name):
+            def hook(module, inp, out):
+                if isinstance(out, torch.Tensor):
+                    v = out.detach().float()
+                    if v.dim() > 1:
+                        v = v.mean(0)
+                    self._activations[name] = v[:32].tolist()
+            return hook
+        for i, layer in enumerate(self.net):
+            layer.register_forward_hook(make_hook(f'net.{i}'))
+
+    def forward(self, x):
+        emb    = self.embedding(x)
+        mask   = (x != 0).float().unsqueeze(-1)
+        pooled = (emb * mask).sum(1) / mask.sum(1).clamp(min=1)
+        return self.net(pooled)
+
+    def get_activations(self) -> dict:
+        return dict(self._activations)
+
+    def get_weight_info(self) -> dict:
+        info = {}
+        for name, param in self.named_parameters():
+            if 'weight' in name and param.dim() == 2:
+                w    = param.detach().float().numpy()
+                r, c = min(w.shape[0], 16), min(w.shape[1], 16)
+                info[name] = {
+                    'shape':    list(w.shape),
+                    'mean_abs': float(np.mean(np.abs(w))),
+                    'std':      float(np.std(w)),
+                    'sample':   w[:r, :c].tolist(),
+                }
+        return info
+
+
+# ─── LIVING NETWORK ───────────────────────────────────────────────────────────
+class LivingNetwork:
+    """
+    The brain. Wraps the PyTorch model with:
+      - knowledge.jsonl  → persistent record of everything it has read
+      - model_checkpoint.pt → saved weights (restored on restart)
+      - training_stats.json → live stats readable by the UI
+    """
+
+    AUTO_SAVE_EVERY = 30   # Save checkpoint every N training epochs
+
+    def __init__(self):
+        self.vocab      = Vocabulary()
+        self.model      = TextClassifier()
+        self.optimizer  = optim.Adam(self.model.parameters(), lr=0.001, weight_decay=1e-5)
+        self.scheduler  = optim.lr_scheduler.ReduceLROnPlateau(
+            self.optimizer, mode='min', patience=20, factor=0.5, min_lr=1e-5)
+        self.criterion  = nn.CrossEntropyLoss()
+
+        self.epoch           = 0
+        self.total_samples   = 0
+        self.data_buffer     = []   # (text, label_int) — in-memory training pool
+        self.loss_history    = []
+        self.acc_history     = []
+        self.category_counts = defaultdict(int)
+        self.knowledge_count = 0    # Total articles ever ingested
+
+        self.stats = {
+            'epoch':         0,
+            'loss':          '—',
+            'accuracy':      '—',
+            'total_samples': 0,
+            'lr':            0.001,
+            'buffer_size':   0,
+            'knowledge_count': 0,
+            'last_text':     '(nothing yet)',
+            'vocab_size':    2,
+            'status':        'idle',
+        }
+
+        # Load checkpoint first, then restore knowledge buffer
+        self._load_checkpoint()
+        self._load_knowledge()
+
+    # ── KNOWLEDGE FILE ────────────────────────────────────────────────────────
+    def _write_knowledge(self, text: str, category: str, source: str = 'unknown'):
+        """Append one learned item to knowledge.jsonl immediately."""
+        record = {
+            'text':      text,
+            'category':  category,
+            'source':    source,
+            'timestamp': datetime.now(UTC).isoformat(),
+            'epoch_at_ingestion': self.epoch,
+        }
+        try:
+            with open(KNOWLEDGE_FILE, 'a', encoding='utf-8') as f:
+                f.write(json.dumps(record, ensure_ascii=False) + '\n')
+            self.knowledge_count += 1
+        except Exception:
+            pass
+
+    def _load_knowledge(self):
+        """On startup: read knowledge.jsonl and rebuild data_buffer + vocab."""
+        if not os.path.exists(KNOWLEDGE_FILE):
+            return
+        loaded = 0
+        try:
+            with open(KNOWLEDGE_FILE, 'r', encoding='utf-8') as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        rec   = json.loads(line)
+                        text  = rec.get('text', '')
+                        cat   = rec.get('category', 'other')
+                        label = CATEGORIES.index(cat) if cat in CATEGORIES else 7
+                        if len(text) > 20:
+                            self.data_buffer.append((text, label))
+                            self.vocab.update(text)
+                            self.category_counts[cat] += 1
+                            loaded += 1
+                    except Exception:
+                        continue
+            self.knowledge_count = loaded
+            if loaded > 0:
+                self.vocab.build()
+            # Trim buffer if huge
+            if len(self.data_buffer) > 5000:
+                self.data_buffer = self.data_buffer[-4000:]
+        except Exception:
+            pass
+        self.stats['knowledge_count'] = self.knowledge_count
+        self.stats['buffer_size']     = len(self.data_buffer)
+        self.stats['vocab_size']      = len(self.vocab)
+
+    def get_knowledge_file_stats(self) -> dict:
+        """Return stats about the knowledge file for the UI."""
+        if not os.path.exists(KNOWLEDGE_FILE):
+            return {'exists': False, 'lines': 0, 'size_kb': 0}
+        size = os.path.getsize(KNOWLEDGE_FILE)
+        lines = 0
+        try:
+            with open(KNOWLEDGE_FILE, 'r', encoding='utf-8') as f:
+                lines = sum(1 for l in f if l.strip())
+        except Exception:
+            pass
+        return {'exists': True, 'lines': lines, 'size_kb': round(size / 1024, 1)}
+
+    def get_recent_knowledge(self, n: int = 20) -> list:
+        """Return last N items from knowledge.jsonl for display."""
+        if not os.path.exists(KNOWLEDGE_FILE):
+            return []
+        lines = []
+        try:
+            with open(KNOWLEDGE_FILE, 'r', encoding='utf-8') as f:
+                all_lines = [l.strip() for l in f if l.strip()]
+            for line in reversed(all_lines[-n:]):
+                try:
+                    lines.append(json.loads(line))
+                except Exception:
+                    pass
+        except Exception:
+            pass
+        return lines
+
+    # ── INGEST ────────────────────────────────────────────────────────────────
+    def ingest(self, text: str, category: str, source: str = 'unknown'):
+        """
+        Add text to training buffer AND write to knowledge.jsonl immediately.
+        This is how the AI 'remembers' what it has learned.
+        """
+        cleaned = text.strip()
+        if len(cleaned) < 20:
+            return
+        label = CATEGORIES.index(category) if category in CATEGORIES else 7
+
+        # ① Write to disk first — never lose this
+        self._write_knowledge(cleaned, category, source)
+
+        # ② Add to in-memory training buffer
+        self.data_buffer.append((cleaned, label))
+        self.vocab.update(cleaned)
+        self.category_counts[category] += 1
+
+        # Rebuild vocab every 25 items
+        if len(self.data_buffer) % 25 == 0:
+            self.vocab.build()
+
+        # Keep buffer bounded (disk has the full history)
+        if len(self.data_buffer) > 5000:
+            self.data_buffer = self.data_buffer[-4000:]
+
+        self.stats.update({
+            'buffer_size':     len(self.data_buffer),
+            'vocab_size':      len(self.vocab),
+            'knowledge_count': self.knowledge_count,
+        })
+
+    # ── TRAINING ──────────────────────────────────────────────────────────────
+    def train_step(self, batch_size: int = 32) -> float | None:
+        if len(self.data_buffer) < batch_size or not self.vocab.is_built:
+            return None
+        with _TEXT_LOCK:
+            self.model.train()
+            idx   = np.random.choice(len(self.data_buffer), batch_size, replace=False)
+            batch = [self.data_buffer[i] for i in idx]
+            texts, labels = zip(*batch)
+
+            x = torch.tensor([self.vocab.encode(t) for t in texts], dtype=torch.long)
+            y = torch.tensor(list(labels), dtype=torch.long)
+
+            self.model.zero_grad(set_to_none=True)
+            logits = self.model(x)
+            loss   = self.criterion(logits, y)
+            loss.backward()
+            for p in self.model.parameters():
+                if p.grad is not None:
+                    p.grad.data.clamp_(-1.0, 1.0)
+            self.optimizer.step()
+            loss_val = loss.detach().item()
+            acc      = (logits.detach().argmax(1) == y).float().mean().item()
+
+        self.epoch        += 1
+        self.total_samples += batch_size
+        self.scheduler.step(loss_val)
+
+        self.loss_history.append(round(loss_val, 5))
+        self.acc_history.append(round(acc, 4))
+        if len(self.loss_history) > 500:
+            self.loss_history = self.loss_history[-500:]
+            self.acc_history  = self.acc_history[-500:]
+
+        self.stats.update({
+            'epoch':           self.epoch,
+            'loss':            round(loss_val, 4),
+            'accuracy':        round(acc * 100, 1),
+            'total_samples':   self.total_samples,
+            'lr':              round(self.optimizer.param_groups[0]['lr'], 7),
+            'buffer_size':     len(self.data_buffer),
+            'last_text':       texts[0][:120],
+            'vocab_size':      len(self.vocab),
+            'knowledge_count': self.knowledge_count,
+        })
+
+        # Auto-save checkpoint every N epochs
+        if self.epoch % self.AUTO_SAVE_EVERY == 0:
+            self.save_checkpoint()
+
+        # Always write stats file so UI can read without waiting
+        self._write_stats_file()
+
+        return loss_val
+
+    def train_n_steps(self, n: int = 50) -> dict:
+        losses = []
+        for _ in range(n):
+            l = self.train_step()
+            if l is not None:
+                losses.append(l)
+        return {
+            'steps':    len(losses),
+            'avg_loss': round(sum(losses) / len(losses), 5) if losses else None,
+        }
+
+    # ── INFERENCE ─────────────────────────────────────────────────────────────
+    def predict(self, text: str) -> dict:
+        if not self.vocab.is_built or not text.strip():
+            return {'error': 'Model not ready — start the network and let it train first'}
+        self.model.eval()
+        with torch.no_grad():
+            x      = torch.tensor([self.vocab.encode(text)], dtype=torch.long)
+            logits = self.model(x)
+            probs  = torch.softmax(logits, dim=1)[0].tolist()
+            pred   = int(logits.argmax(1).item())
+        return {
+            'prediction': CATEGORIES[pred],
+            'confidence': round(probs[pred] * 100, 1),
+            'all_probs':  {c: round(p * 100, 2) for c, p in zip(CATEGORIES, probs)},
+        }
+
+    # ── VIZ STATE ─────────────────────────────────────────────────────────────
+    def get_viz_state(self) -> dict:
+        self.model.eval()
+        with torch.no_grad():
+            dummy = torch.zeros(1, 64, dtype=torch.long)
+            self.model(dummy)
+        return {
+            'layer_sizes':    [self.model.embed_dim] + self.model.hidden_dims + [self.model.num_classes],
+            'activations':    self.model.get_activations(),
+            'weights':        self.model.get_weight_info(),
+            'loss_history':   self.loss_history[-100:],
+            'acc_history':    self.acc_history[-100:],
+            'stats':          self.stats,
+            'category_counts':dict(self.category_counts),
+        }
+
+    # ── PERSISTENCE ───────────────────────────────────────────────────────────
+    def save_checkpoint(self):
+        try:
+            torch.save({
+                'model':           self.model.state_dict(),
+                'optimizer':       self.optimizer.state_dict(),
+                'epoch':           self.epoch,
+                'total_samples':   self.total_samples,
+                'loss_history':    self.loss_history,
+                'acc_history':     self.acc_history,
+                'vocab_word2idx':  self.vocab.word2idx,
+                'category_counts': dict(self.category_counts),
+                'stats':           self.stats,
+            }, CHECKPOINT_FILE)
+            return True
+        except Exception:
+            return False
+
+    def _load_checkpoint(self):
+        if not os.path.exists(CHECKPOINT_FILE):
+            return False
+        try:
+            ck = torch.load(CHECKPOINT_FILE, map_location='cpu')
+            self.model.load_state_dict(ck['model'])
+            self.optimizer.load_state_dict(ck['optimizer'])
+            self.epoch           = ck.get('epoch', 0)
+            self.total_samples   = ck.get('total_samples', 0)
+            self.loss_history    = ck.get('loss_history', [])
+            self.acc_history     = ck.get('acc_history', [])
+            self.category_counts = defaultdict(int, ck.get('category_counts', {}))
+            self.stats           = ck.get('stats', self.stats)
+            w2i = ck.get('vocab_word2idx', {})
+            if w2i:
+                self.vocab.word2idx  = w2i
+                self.vocab.idx2word  = {v: k for k, v in w2i.items()}
+                self.vocab.is_built  = len(w2i) > 2
+            return True
+        except Exception:
+            return False
+
+    def _write_stats_file(self):
+        """Write human-readable stats to training_stats.json for easy debugging."""
+        try:
+            with open(STATS_FILE, 'w') as f:
+                json.dump({**self.stats, 'loss_last10': self.loss_history[-10:]}, f, indent=2)
+        except Exception:
+            pass