"""Generate a short English continuation with Custom GPT 40M.
Copyright (c) 2026 sraivante. SPDX-License-Identifier: Apache-2.0
"""
import argparse
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model",default="sraivante/Custom-GPT-40M-Base")
    parser.add_argument("--revision",default="main",help="Pin a Hub commit for reproducibility.")
    parser.add_argument("--prompt",default="The future of artificial intelligence")
    parser.add_argument("--max-new-tokens",type=int,default=80)
    parser.add_argument("--seed",type=int,default=42)
    parser.add_argument("--greedy",action="store_true")
    parser.add_argument("--device",default="cpu",choices=["cpu","cuda"])
    args=parser.parse_args()
    if not args.prompt.strip():parser.error("Provide a non-empty English prompt.")
    if args.max_new_tokens<1:parser.error("--max-new-tokens must be positive.")
    tokenizer=AutoTokenizer.from_pretrained(args.model,revision=args.revision)
    model=AutoModelForCausalLM.from_pretrained(args.model,revision=args.revision).to(args.device).eval()
    inputs=tokenizer(args.prompt,return_tensors="pt",add_special_tokens=False).to(args.device)
    remaining=model.config.n_positions-inputs["input_ids"].shape[1]
    if remaining<1:parser.error("The prompt must be shorter than 256 tokens.")
    torch.manual_seed(args.seed)
    options={} if args.greedy else {"temperature":0.8,"top_k":50}
    with torch.inference_mode():
        output=model.generate(**inputs,max_new_tokens=min(args.max_new_tokens,remaining),do_sample=not args.greedy,**options)
    print(tokenizer.decode(output[0],skip_special_tokens=True))
if __name__=="__main__":main()