Bugfixes for KBLaM

[ti250]

Update documentation, training scripts, and utility functions to improve usability, clarity, and compatibility with additional configurations. Key changes include enhancements to the README.md for better guidance, updates to the training script for flexibility and improved defaults, and modifications to utility functions to support additional arguments and models.

Credits to @xidulu for most of the changes (from his [gist] (https://gist.github.com/xidulu/dec7a154179a552c9a24856a52412491#file-kblam-diff-L93-L94)).