| """Generate text from alpha-er with the PyTorch port. |
| |
| Pads to block_size and reads the last real position, which the conditional MLP |
| requires and which is exact: attention is causal, and each token's expert is a |
| function of its own position. |
| """ |
| import sys, os, json, torch |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from modeling_alpha import AlphaErConfig, AlphaErForCausalLM |
| from tokenization_alpha import AlphaErTokenizer |
| from safetensors.torch import load_file |
|
|
| hf_dir = sys.argv[1] |
| prompts = sys.argv[2:] or ["<|user|>Hello!<|assistant|>"] |
| temp = float(os.environ.get("TEMP", "0.8")) |
| topk = int(os.environ.get("TOPK", "40")) |
| ntok = int(os.environ.get("NTOK", "60")) |
|
|
| cfg_d = json.load(open(f"{hf_dir}/config.json")) |
| cfg = AlphaErConfig(**{k: v for k, v in cfg_d.items() |
| if k in AlphaErConfig.__init__.__code__.co_varnames}) |
| model = AlphaErForCausalLM(cfg) |
| model.load_state_dict(load_file(f"{hf_dir}/model.safetensors"), strict=False) |
| model.eval() |
| tok = AlphaErTokenizer.from_file(f"{hf_dir}/tokenizer_artifacts.json") |
| torch.manual_seed(1234) |
|
|
| print(f"alpha-er step {cfg_d.get('trained_step')} temp={temp} top_k={topk}\n") |
| for p in prompts: |
| ids = tok.encode(p) |
| out = model.generate(torch.tensor([ids]), max_new_tokens=ntok, |
| temperature=temp, top_k=topk)[0].tolist() |
| print("=" * 72) |
| print("PROMPT:", p) |
| print("OUTPUT:", tok.decode(out[len(ids):]).replace("\n", "\\n")) |
|
|
