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