Upload generate.py with huggingface_hub

generate.py

@@ -0,0 +1,85 @@
+"""
+Eve-2-MoE Inference
+===================
+Quick generation script. Works with local weights or HuggingFace download.
+
+Usage:
+    python generate.py --prompt "The future of AI is"
+    python generate.py --prompt "The future of AI is" --model_path ./model_final/pytorch_model.bin
+    python generate.py --prompt "The future of AI is" --hf_repo anthonym21/Eve-2-MoE-250M
+"""
+
+import argparse
+import torch
+import tiktoken
+from modeling_eve import ModelConfig, DeepSeekMoE
+
+
+def load_model(model_path: str = None, hf_repo: str = None, device: str = "cuda"):
+    config = ModelConfig()
+    model = DeepSeekMoE(config)
+
+    if hf_repo:
+        from huggingface_hub import hf_hub_download
+        model_path = hf_hub_download(repo_id=hf_repo, filename="pytorch_model.bin")
+
+    if model_path:
+        state_dict = torch.load(model_path, map_location=device, weights_only=True)
+        model.load_state_dict(state_dict)
+
+    return model.to(device).eval()
+
+
+def generate_streaming(model, prompt: str, max_tokens: int = 200,
+                       temperature: float = 0.8, top_k: int = 50, device: str = "cuda"):
+    enc = tiktoken.get_encoding("gpt2")
+    tokens = torch.tensor(enc.encode(prompt), dtype=torch.long, device=device).unsqueeze(0)
+
+    print(prompt, end="", flush=True)
+
+    with torch.no_grad():
+        for _ in range(max_tokens):
+            idx_cond = tokens[:, -model.config.block_size:]
+
+            with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=(device == "cuda")):
+                logits, _ = model(idx_cond)
+
+            logits = logits[:, -1, :] / temperature
+
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float("Inf")
+
+            probs = torch.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            tokens = torch.cat((tokens, idx_next), dim=1)
+
+            print(enc.decode([idx_next.item()]), end="", flush=True)
+
+    print("\n")
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--prompt", type=str, default="The future of artificial intelligence is")
+    p.add_argument("--model_path", type=str, default=None)
+    p.add_argument("--hf_repo", type=str, default=None)
+    p.add_argument("--max_tokens", type=int, default=200)
+    p.add_argument("--temperature", type=float, default=0.8)
+    p.add_argument("--top_k", type=int, default=50)
+    p.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu")
+    args = p.parse_args()
+
+    if not args.model_path and not args.hf_repo:
+        args.hf_repo = "anthonym21/Eve-2-MoE-250M"
+
+    print(f"Loading model on {args.device}...")
+    model = load_model(args.model_path, args.hf_repo, args.device)
+    param_count = sum(p.numel() for p in model.parameters())
+    print(f"Parameters: {param_count / 1e6:.2f}M\n")
+
+    generate_streaming(model, args.prompt, args.max_tokens, args.temperature, args.top_k, args.device)
+
+
+if __name__ == "__main__":
+    main()