File size: 2,741 Bytes
2b66baf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | """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))
|