"""Real inference using Astheral's merged weights. Run on a remote GPU."""
import argparse,json,torch
from pathlib import Path
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer,AutoModelForCausalLM,BitsAndBytesConfig,set_seed
p=argparse.ArgumentParser()
p.add_argument("prompt")
p.add_argument("--mode",choices=["instruct","thinking"],default="instruct")
p.add_argument("--model",default="SEE7/Astheral")
p.add_argument("--system",default=None,help="Your own system instructions; replaces the default assistant instructions.")
p.add_argument("--max-new-tokens",type=int,default=1024)
p.add_argument("--load-in-4bit",action="store_true",help="Load weights in NF4 to reduce remote GPU memory use.")
p.add_argument("--seed",type=int,default=20261113)
args=p.parse_args()
if not torch.cuda.is_available():raise RuntimeError("GPU required; no CPU fallback.")
set_seed(args.seed)
settings_path=Path(args.model)/"settings.json" if Path(args.model).is_dir() else Path(hf_hub_download(args.model,"settings.json"))
settings=json.loads(settings_path.read_text())
system=(args.system if args.system is not None else settings["system"])+("\n"+settings["thinking"] if args.mode=="thinking" else "")
tok=AutoTokenizer.from_pretrained(args.model)
dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
quant=BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type="nf4",bnb_4bit_use_double_quant=True,bnb_4bit_compute_dtype=dtype) if args.load_in_4bit else None
model=AutoModelForCausalLM.from_pretrained(args.model,torch_dtype=dtype,device_map={"":"cuda:0"},quantization_config=quant)
inputs=tok.apply_chat_template([{"role":"system","content":system},{"role":"user","content":args.prompt}],tokenize=True,add_generation_prompt=True,return_dict=True,return_tensors="pt").to(model.device)
with torch.inference_mode():
    output=model.generate(**inputs,max_new_tokens=args.max_new_tokens,do_sample=True,temperature=0.6 if args.mode=="thinking" else 0.7,top_p=0.95 if args.mode=="thinking" else 0.8,top_k=20,pad_token_id=128004,eos_token_id=128009)
print(tok.decode(output[0,inputs["input_ids"].shape[1]:],skip_special_tokens=True))