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