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