File size: 1,464 Bytes
d7562c8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | """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"))
|