| """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 AutoConfig,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) |
| p.add_argument("--sample",action="store_true",help="Enable sampling for creative variation; defaults to deterministic decoding.") |
| 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 |
| config=AutoConfig.from_pretrained(args.model) |
| saved_quant=getattr(config,"quantization_config",None) |
| if quant is None and isinstance(saved_quant,dict) and saved_quant.get("load_in_4bit"): |
| quant=BitsAndBytesConfig.from_dict(saved_quant) |
| quant.bnb_4bit_compute_dtype=dtype |
| config.quantization_config=quant.to_dict() |
| model=AutoModelForCausalLM.from_pretrained(args.model,config=config,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(): |
| generation={"max_new_tokens":args.max_new_tokens,"do_sample":args.sample,"pad_token_id":128004,"eos_token_id":128009} |
| if args.sample: |
| generation.update(temperature=0.6 if args.mode=="thinking" else 0.7,top_p=0.95 if args.mode=="thinking" else 0.8,top_k=20) |
| output=model.generate(**inputs,**generation) |
| print(tok.decode(output[0,inputs["input_ids"].shape[1]:],skip_special_tokens=True)) |
|
|
