File size: 2,867 Bytes
4cf8d9d | 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 39 | """Run Bharat24B with its saved adapter and native tokenizer on a CUDA GPU."""
import argparse,hashlib,json
from pathlib import Path
def main():
import torch
from huggingface_hub import snapshot_download
from peft import PeftModel
from transformers import BitsAndBytesConfig,Mistral3ForConditionalGeneration,MistralCommonBackend
from bharat_serving import encode_messages,visible_answer
p=argparse.ArgumentParser(description=__doc__)
p.add_argument('--prompt',required=True);p.add_argument('--revision',default=None)
p.add_argument('--max-new-tokens',type=int,default=4096)
p.add_argument('--no-identity-policy',action='store_true');p.add_argument('--raw-output',action='store_true')
args=p.parse_args();assert args.max_new_tokens>0 and torch.cuda.is_available(),'A CUDA GPU is required.'
folder=Path(snapshot_download('DesaiJayM/Bharat-24B-Math-Code',revision=args.revision,
allow_patterns=['adapter_model.safetensors','adapter_config.json','tekken.json','release-manifest.json','bharat-model-identity.json']))
manifest=json.loads((folder/'release-manifest.json').read_text())
for name,digest in manifest['unchanged_checkpoint_files_sha256'].items():
assert hashlib.sha256((folder/name).read_bytes()).hexdigest()==digest,name
tokenizer=MistralCommonBackend.from_pretrained(folder,mode='test')
backbone=Mistral3ForConditionalGeneration.from_pretrained(manifest['base_model'],revision=manifest['base_revision'],
quantization_config=BitsAndBytesConfig(load_in_4bit=True,bnb_4bit_quant_type='nf4',
bnb_4bit_use_double_quant=True,bnb_4bit_compute_dtype=torch.bfloat16),
device_map={'':0},dtype=torch.bfloat16,attn_implementation='sdpa',low_cpu_mem_usage=True)
model=PeftModel.from_pretrained(backbone,folder,is_trainable=False).eval();model.config.use_cache=True
messages=[{'role':'user','content':args.prompt}]
ids=tokenizer.apply_chat_template(messages,tokenize=True,return_dict=False,add_generation_prompt=True) if args.no_identity_policy else encode_messages(tokenizer,messages,folder/'bharat-model-identity.json')
inputs=torch.tensor([ids],dtype=torch.long,device='cuda')
with torch.inference_mode(),torch.autocast('cuda',dtype=torch.bfloat16):
output=model.generate(input_ids=inputs,attention_mask=torch.ones_like(inputs),do_sample=False,
max_new_tokens=args.max_new_tokens,use_cache=True,pad_token_id=tokenizer.pad_token_id or 11,eos_token_id=tokenizer.eos_token_id)
tokens=output[0,len(ids):].tolist();raw=tokenizer.decode(tokens,skip_special_tokens=True)
print(visible_answer(raw))
if args.raw_output:print('\nFull response:\n'+raw)
print('\nGenerated tokens:',len(tokens),'Reached token cap:',len(tokens)>=args.max_new_tokens and tokenizer.eos_token_id not in tokens)
if __name__=='__main__':main()
|