address bot's comments

Author: mtake

examples/README.md

@@ -112,14 +112,14 @@ training_hub has a utility for merging two checkpoints of the same model into on
 
 **Command-Line Example:**
 ```bash
-python interpolator.py --model-path ibm-granite/granite-3.3-8b-instruct --trained-model-path /path/to/checkpoint
+python interpolator.py --model-path /path/to/base/model --trained-model-path /path/to/trained/checkpoint
 ```
 
 **Python Example:**
 ```python
 from interpolator import interpolate_models
 
-interpolate_models("ibm-granite/granite-3.3-8b-instruct", "/path/to/checkpoint")
+interpolate_models("/path/to/base/model", "/path/to/trained/checkpoint")
 ```
 
 ## Getting Started

examples/scripts/interpolator.py

@@ -5,8 +5,8 @@
 
 Example usage:
     python interpolator.py \\
-        --model-path ibm-granite/granite-3.3-8b-instruct \\
-        --trained-model-path /path/to/checkpoint
+        --model-path /path/to/base/model \\
+        --trained-model-path /path/to/trained/checkpoint
 """
 # Standard
 import argparse
@@ -47,15 +47,15 @@ def interpolate_models(
         else:
             model_kwargs["torch_dtype"] = torch_dtype
 
-    # load original model
+    # load base model
     model = AutoModelForCausalLM.from_pretrained(
         model_path,
         **model_kwargs,
     )
     state_dict = model.state_dict()
-    original_model_weight = 1 - trained_model_weight
+    base_model_weight = 1 - trained_model_weight
     for key in state_dict.keys():
-        state_dict[key] = state_dict[key] * original_model_weight
+        state_dict[key] = state_dict[key] * base_model_weight
 
     # load trained model
     trained_model = AutoModelForCausalLM.from_pretrained(
@@ -66,13 +66,15 @@ def interpolate_models(
     for key in state_dict.keys():
         state_dict[key] += trained_state_dict[key] * trained_model_weight
 
-    # save interpolated model
+    # save merged model
     model.save_pretrained(output_model_path, state_dict=state_dict)
 
     # copy tokenizer
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     tokenizer.save_pretrained(output_model_path)
 
+    print(f"Merged model saved at {output_model_path}")
+
     return output_model_path
 
 
@@ -84,7 +86,7 @@ def parse_arguments():
         "--model-path",
         type=str,
         required=True,
-        help="Path to the original model",
+        help="Path to the base model",
     )
     parser.add_argument(
         "--trained-model-path",

examples/scripts/osft_granite_example.py

@@ -120,7 +120,7 @@ def main():
     parser.add_argument('--learning-rate', type=float, default=default_learning_rate,
                        help=f'Learning rate for training (default: {default_learning_rate})')
     parser.add_argument('--unmask-messages', action='store_true', default=False,
-                       help='Unmask messages during training (default: False)')
+                       help='Unmask all non-system messages during training, otherwise only unmasks assistant messages (default: False)')
     parser.add_argument('--batch-size', type=int, default=default_batch_size,
                        help=f'Effective batch size for training (default: {default_batch_size})')
     parser.add_argument('--max-seq-len', type=int, default=default_max_seq_len,