Upload chat_infer.py with huggingface_hub

chat_infer.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+"""Fast CogNet inference for chat API. Called by Next.js API route."""
+import sys, os, json, time
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+import torch
+from cognet_1b import CogNet1B
+from infer import CharTokenizer
+
+# Globals - load once, reuse across calls
+_model = None
+_tokenizer = None
+
+def load_model():
+    global _model, _tokenizer
+    if _model is not None:
+        return _model, _tokenizer
+    
+    ckpt_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'checkpoints')
+    _tokenizer = CharTokenizer.load(os.path.join(ckpt_dir, 'tokenizer_v3.json'))
+    _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
+    )
+    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']
+        # Handle FP16 weights
+        fp16_state = {k: v.float() if v.dtype == torch.float16 else v for k, v in state.items()}
+        _model.load_state_dict(fp16_state)
+        step = ckpt.get('metrics', {}).get('step', '?')
+        sys.stderr.write(f'Model loaded (step={step})\n')
+    else:
+        sys.stderr.write('WARNING: No checkpoint found, using random weights\n')
+    _model.eval()
+    return _model, _tokenizer
+
+def main():
+    prompt = sys.argv[1] if len(sys.argv) > 1 else "Hello"
+    max_tokens = int(sys.argv[2]) if len(sys.argv) > 2 else 80
+    temperature = float(sys.argv[3]) if len(sys.argv) > 3 else 0.7
+    top_k = int(sys.argv[4]) if len(sys.argv) > 4 else 30
+
+    model, tokenizer = load_model()
+    
+    ids = tokenizer.encode(prompt)
+    if not ids:
+        ids = [0]
+    input_ids = torch.tensor([ids], dtype=torch.long)
+    
+    t0 = time.time()
+    with torch.no_grad():
+        gen = model.generate(input_ids, max_new_tokens=max_tokens, temperature=temperature, top_k=top_k)
+    elapsed = time.time() - t0
+    
+    generated_ids = gen[0].tolist()
+    generated_text = tokenizer.decode(generated_ids)
+    new_text = tokenizer.decode(generated_ids[len(ids):])
+    
+    result = {
+        'generated_text': generated_text,
+        'new_text': new_text,
+        'num_tokens': len(generated_ids),
+        'inference_time_ms': round(elapsed * 1000),
+        'prompt': prompt,
+    }
+    print(json.dumps(result, ensure_ascii=False))
+
+if __name__ == '__main__':
+    main()