Upload infer.py with huggingface_hub

infer.py

@@ -0,0 +1,385 @@
+"""
+CogNet Inference Engine for Next.js API
+=======================================
+Loads trained CogNet model and CharTokenizer, supports:
+- generate: text generation with temperature/top-k sampling
+- analyze: logits analysis, entropy, top predictions
+- inspect: model architecture details
+- info: model info without loading weights
+"""
+
+import json
+import math
+import os
+import sys
+from typing import Any, Dict, List, Optional
+
+import torch
+import torch.nn.functional as F
+
+# Import from same directory
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from cognet_1b import CogNet1B
+
+
+# ─── Model Config (matches training) ────────────────────────────────────────
+
+MODEL_CONFIG = {
+    'vocab_size': 136,
+    '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,
+}
+
+CKPT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'checkpoints')
+TOKENIZER_PATH = os.path.join(CKPT_DIR, 'tokenizer_v3.json')
+BEST_MODEL_PATH = os.path.join(CKPT_DIR, 'cognet_best.pt')
+LATEST_MODEL_PATH = os.path.join(CKPT_DIR, 'cognet_latest.pt')
+
+
+# ─── CharTokenizer (standalone, no import needed from train_pipeline) ───────
+
+class CharTokenizer:
+    """Character-level tokenizer: printable ASCII + French accents + newline/tab."""
+
+    def __init__(self):
+        self.chars = sorted(set(
+            [chr(i) for i in range(32, 127)]
+            + list('àâäéèêëïîôùûüÿçœæÀÂÄÉÈÊËÏÎÔÙÛÜŸÇŒÆ')
+            + list('ëßñ¿«»')
+            + ['\t', '\n']
+        ))
+        self.char_to_id = {c: i for i, c in enumerate(self.chars)}
+        self.id_to_char = {i: c for i, c in enumerate(self.chars)}
+        self.vocab_size = len(self.chars)
+
+    def encode(self, text: str) -> List[int]:
+        return [self.char_to_id.get(c, self.char_to_id.get(' ', 0)) for c in text]
+
+    def decode(self, ids: List[int]) -> str:
+        return ''.join(self.id_to_char.get(i, ' ') for i in ids)
+
+    def save(self, path: str):
+        with open(path, 'w', encoding='utf-8') as f:
+            json.dump({
+                'chars': self.chars,
+                'vocab_size': self.vocab_size,
+            }, f, ensure_ascii=False, indent=2)
+
+    @classmethod
+    def load(cls, path: str) -> 'CharTokenizer':
+        tok = cls.__new__(cls)
+        with open(path, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        tok.chars = data['chars']
+        tok.char_to_id = {c: i for i, c in enumerate(tok.chars)}
+        tok.id_to_char = {i: c for i, c in enumerate(tok.chars)}
+        tok.vocab_size = data['vocab_size']
+        return tok
+
+
+# ─── JSON Helpers ────────────────────────────────────────────────────────────
+
+def sanitize_for_json(obj: Any) -> Any:
+    """Replace NaN/Inf with None for JSON serialization."""
+    if isinstance(obj, float):
+        if math.isnan(obj) or math.isinf(obj):
+            return None
+        return obj
+    if isinstance(obj, dict):
+        return {k: sanitize_for_json(v) for k, v in obj.items()}
+    if isinstance(obj, list):
+        return [sanitize_for_json(v) for v in obj]
+    return obj
+
+
+# ─── Model Cache ─────────────────────────────────────────────────────────────
+
+_model_cache: Dict[str, Any] = {
+    'model': None,
+    'tokenizer': None,
+    'device': None,
+    'loaded': False,
+}
+
+
+def load_model_and_tokenizer() -> tuple:
+    """Load model and tokenizer with caching."""
+    if _model_cache['loaded']:
+        return _model_cache['model'], _model_cache['tokenizer'], _model_cache['device']
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    # Load tokenizer
+    if not os.path.exists(TOKENIZER_PATH):
+        raise FileNotFoundError(
+            f"Tokenizer not found at {TOKENIZER_PATH}. "
+            "Run train_pipeline.py first to create it."
+        )
+    tokenizer = CharTokenizer.load(TOKENIZER_PATH)
+
+    # Update vocab_size from tokenizer
+    config = dict(MODEL_CONFIG)
+    config['vocab_size'] = tokenizer.vocab_size
+
+    # Create model
+    model = CogNet1B(**config).to(device)
+
+    # Load weights (prefer best, then latest)
+    model_path = BEST_MODEL_PATH if os.path.exists(BEST_MODEL_PATH) else LATEST_MODEL_PATH
+    if model_path and os.path.exists(model_path):
+        ckpt = torch.load(model_path, map_location=device, weights_only=False)
+        model.load_state_dict(ckpt['model_state_dict'])
+        step = ckpt.get('metrics', {}).get('step', '?')
+        print(f"Loaded model from {model_path} (step={step})")
+    else:
+        print("WARNING: No trained weights found. Using random initialization.")
+
+    model.eval()
+
+    # Cache
+    _model_cache['model'] = model
+    _model_cache['tokenizer'] = tokenizer
+    _model_cache['device'] = device
+    _model_cache['loaded'] = True
+
+    return model, tokenizer, device
+
+
+# ─── Action Handlers ─────────────────────────────────────────────────────────
+
+def handle_generate(prompt: str, max_tokens: int = 100,
+                    temperature: float = 0.8, top_k: int = 20) -> Dict:
+    """Generate text from a prompt."""
+    model, tokenizer, device = load_model_and_tokenizer()
+
+    # Encode prompt
+    ids = tokenizer.encode(prompt)
+    if len(ids) == 0:
+        ids = [0]
+
+    input_ids = torch.tensor([ids], dtype=torch.long, device=device)
+
+    # Generate
+    with torch.no_grad():
+        output_ids = model.generate(
+            input_ids,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            top_k=top_k,
+        )
+
+    # Decode
+    generated_ids = output_ids[0].tolist()
+    generated_text = tokenizer.decode(generated_ids)
+    new_text = tokenizer.decode(generated_ids[len(ids):])
+
+    # Token details
+    token_details = []
+    for i, tid in enumerate(generated_ids):
+        char = tokenizer.decode([tid])
+        token_details.append({
+            'id': tid,
+            'char': char,
+            'position': i,
+        })
+
+    return sanitize_for_json({
+        'action': 'generate',
+        'prompt': prompt,
+        'generated_text': generated_text,
+        'new_text': new_text,
+        'token_details': token_details,
+        'num_tokens': len(generated_ids),
+        'temperature': temperature,
+        'top_k': top_k,
+    })
+
+
+def handle_analyze(prompt: str) -> Dict:
+    """Analyze logits, entropy, and top predictions."""
+    model, tokenizer, device = load_model_and_tokenizer()
+
+    ids = tokenizer.encode(prompt)
+    if len(ids) == 0:
+        ids = [0]
+
+    input_ids = torch.tensor([ids], dtype=torch.long, device=device)
+
+    with torch.no_grad():
+        result = model(input_ids, return_stats=True)
+        logits = result['logits']
+
+    # Analyze last token's predictions
+    last_logits = logits[0, -1, :]  # (vocab_size,)
+    probs = F.softmax(last_logits, dim=-1)
+
+    # Entropy
+    entropy = -(probs * (probs + 1e-10).log()).sum().item()
+
+    # Top 10 predictions
+    topk_vals, topk_ids = torch.topk(probs, min(10, probs.size(0)))
+    top_predictions = []
+    for prob, tid in zip(topk_vals.tolist(), topk_ids.tolist()):
+        top_predictions.append({
+            'token_id': tid,
+            'char': tokenizer.decode([tid]),
+            'probability': prob,
+        })
+
+    # Per-position entropy
+    all_probs = F.softmax(logits[0], dim=-1)
+    pos_entropy = (-(all_probs * (all_probs + 1e-10).log()).sum(dim=-1)).tolist()
+
+    # Stats
+    stats = result.get('stats', {})
+    stats_summary = {}
+    for k, v in stats.items():
+        if isinstance(v, torch.Tensor):
+            v = v.item()
+        if isinstance(v, float) and (math.isnan(v) or math.isinf(v)):
+            v = None
+        stats_summary[k] = v
+
+    return sanitize_for_json({
+        'action': 'analyze',
+        'prompt': prompt,
+        'prompt_length': len(ids),
+        'entropy': entropy,
+        'top_predictions': top_predictions,
+        'per_position_entropy': pos_entropy,
+        'model_stats': stats_summary,
+    })
+
+
+def handle_inspect() -> Dict:
+    """Return model architecture details."""
+    model, tokenizer, device = load_model_and_tokenizer()
+
+    params = model.count_parameters()
+    complexity = model.get_complexity_analysis()
+
+    # Layer details
+    layers = []
+    for i, block in enumerate(model.blocks):
+        layer_params = sum(p.numel() for p in block.parameters())
+        layers.append({
+            'block_index': i,
+            'parameters': layer_params,
+            'components': ['CognitiveRouter', 'SharedHierarchicalMemory',
+                          'AdaptiveComputationBlock', 'CompositionalReasoner'],
+        })
+
+    return sanitize_for_json({
+        'action': 'inspect',
+        'architecture': 'CogNet (Non-Transformer)',
+        'total_parameters': params['total'],
+        'trainable_parameters': params['trainable'],
+        'config': {
+            'vocab_size': model.vocab_size,
+            'hidden_dim': model.hidden_dim,
+            'num_blocks': model.num_blocks,
+            'num_channels': model.num_channels,
+            'channel_dim': model.channel_dim,
+            'ff_dim': model.ff_dim,
+            'max_seq_len': model.max_seq_len,
+            'tokenizer_vocab_size': tokenizer.vocab_size,
+        },
+        'complexity_analysis': complexity,
+        'layers': layers,
+        'device': str(device),
+    })
+
+
+def handle_info() -> Dict:
+    """Return model info without loading weights."""
+    config = dict(MODEL_CONFIG)
+
+    # Check what's available
+    has_tokenizer = os.path.exists(TOKENIZER_PATH)
+    has_best = os.path.exists(BEST_MODEL_PATH)
+    has_latest = os.path.exists(LATEST_MODEL_PATH)
+
+    # Estimate param count without loading
+    model = CogNet1B(**config)
+    params = model.count_parameters()
+
+    # Check checkpoint info if available
+    checkpoint_info = {}
+    if has_best:
+        try:
+            ckpt = torch.load(BEST_MODEL_PATH, map_location='cpu', weights_only=False)
+            checkpoint_info['best'] = {
+                'step': ckpt.get('metrics', {}).get('step', None),
+                'val_loss': ckpt.get('metrics', {}).get('val_loss', None),
+                'val_ppl': ckpt.get('metrics', {}).get('val_ppl', None),
+            }
+        except Exception:
+            checkpoint_info['best'] = {'error': 'Could not read checkpoint'}
+    if has_latest:
+        try:
+            ckpt = torch.load(LATEST_MODEL_PATH, map_location='cpu', weights_only=False)
+            checkpoint_info['latest'] = {
+                'step': ckpt.get('metrics', {}).get('step', None),
+            }
+        except Exception:
+            checkpoint_info['latest'] = {'error': 'Could not read checkpoint'}
+
+    return sanitize_for_json({
+        'action': 'info',
+        'model_name': 'CogNet',
+        'architecture': 'Non-Transformer (Cognitive Routing)',
+        'estimated_parameters': params['total'],
+        'config': config,
+        'files': {
+            'tokenizer': has_tokenizer,
+            'best_checkpoint': has_best,
+            'latest_checkpoint': has_latest,
+        },
+        'checkpoint_info': checkpoint_info,
+    })
+
+
+# ─── CLI Entry Point ─────────────────────────────────────────────────────────
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser(description='CogNet Inference Engine')
+    parser.add_argument('action', choices=['generate', 'analyze', 'inspect', 'info'],
+                        help='Action to perform')
+    parser.add_argument('--prompt', type=str, default='The ',
+                        help='Prompt text (for generate/analyze)')
+    parser.add_argument('--max-tokens', type=int, default=100,
+                        help='Max tokens to generate')
+    parser.add_argument('--temperature', type=float, default=0.8,
+                        help='Sampling temperature')
+    parser.add_argument('--top-k', type=int, default=20,
+                        help='Top-k sampling')
+
+    args = parser.parse_args()
+
+    if args.action == 'generate':
+        result = handle_generate(args.prompt, args.max_tokens,
+                                args.temperature, args.top_k)
+    elif args.action == 'analyze':
+        result = handle_analyze(args.prompt)
+    elif args.action == 'inspect':
+        result = handle_inspect()
+    elif args.action == 'info':
+        result = handle_info()
+
+    print(json.dumps(result, indent=2, ensure_ascii=False))
+
+
+if __name__ == '__main__':
+    main()