"""LLaMA-Factory entry with a read-only quantization/memory evidence callback."""
import json,sys,torch,collections
from pathlib import Path
from transformers import TrainerCallback
from llamafactory.train.tuner import run_exp
class Evidence(TrainerCallback):
 def on_train_begin(self,args,state,control,model=None,**kwargs):
  quant=getattr(model.config,'quantization_config',None)
  if hasattr(quant,'to_dict'):quant=quant.to_dict()
  data={'is_loaded_in_4bit':bool(getattr(model,'is_loaded_in_4bit',False)),'quantization_config':quant,'linear4bit_layers':sum(type(m).__name__=='Linear4bit' for m in model.modules()),'linear4bit_types':dict(collections.Counter(type(m).__module__+'.'+type(m).__name__ for m in model.modules() if type(m).__name__=='Linear4bit')),'trainable_parameter_dtypes':sorted({str(p.dtype) for p in model.parameters() if p.requires_grad})}
  assert data['is_loaded_in_4bit'] and data['linear4bit_layers']>0, data
  assert quant['bnb_4bit_quant_type']=='nf4' and quant['bnb_4bit_use_double_quant'] is True, data
  (Path(args.output_dir)/'quantization_runtime.json').write_text(json.dumps(data,indent=2));print('QUANTIZATION_RUNTIME '+json.dumps(data),flush=True)
run_exp(callbacks=[Evidence()])