| """ | |
| 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() | |