Great work! I noticed that you initialized the LWM model with LLaMA-2 7B model. But I couldn't find where you loaded it anywhere (certainly not in the scripts/run_train_text.sh). Would you tell me which model weight you loaded and how to load it for the Stage I LLM tuning? Thanks!