File size: 2,252 Bytes
3b85381 84478eb 3b85381 84478eb 3b85381 | 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 37 38 39 40 41 42 43 44 | """Cached greedy completion for the custom AMOR release."""
import argparse
from contextlib import nullcontext
import json
from pathlib import Path
import torch
from transformers import AutoTokenizer
from load_model import load_model
from amor_decode import AMORDecodeWrapper
def main():
p=argparse.ArgumentParser(description=__doc__)
p.add_argument("--router", action="store_true", help="Use the fitted inference routers.")
p.add_argument("--prompt",required=True)
p.add_argument("--max-new-tokens",type=int,default=32)
p.add_argument("--device",choices=["cuda","cpu"],default="cuda")
p.add_argument("--model-dir",type=Path,default=Path(__file__).resolve().parent)
a=p.parse_args()
if a.max_new_tokens<1:p.error("--max-new-tokens must be positive")
config=json.loads((a.model_dir/"config.json").read_text())
if a.device=="cuda" and not torch.cuda.is_available():
p.error("CUDA is unavailable; install a CUDA-enabled PyTorch build, or use --device cpu for Mamba2.")
if a.device=="cpu" and config["model_kwargs"]["backbone"]=="gdn":
p.error("This Gated DeltaNet implementation requires compatible GPU kernels.")
tokenizer=AutoTokenizer.from_pretrained(config["tokenizer"],revision=config.get("tokenizer_revision"))
model=load_model(a.model_dir,device=a.device,dtype=torch.float32,use_router=a.router)
ids=tokenizer(a.prompt,return_tensors="pt")["input_ids"].to(a.device)
if not ids.shape[1]:p.error("Prompt must tokenize to at least one token")
wrapper=AMORDecodeWrapper(model)
cast=torch.autocast("cuda",dtype=torch.bfloat16) if a.device=="cuda" else nullcontext()
generated=[]
with torch.inference_mode(),cast:
logits,cache=wrapper.prefill(ids)
for i in range(a.max_new_tokens):
scores=logits[:,-1] if logits.ndim==3 else logits
token=scores.argmax(-1,keepdim=True)
value=int(token.item())
if value==tokenizer.eos_token_id:break
generated.append(value)
if i+1<a.max_new_tokens:
logits,cache,_=wrapper.decode_step(token,cache,step_idx=ids.shape[1]+i)
print(a.prompt+tokenizer.decode(generated,skip_special_tokens=True))
if __name__=="__main__":main()
|