Address bot comments

Author: mtake

examples/scripts/interpolator.py

@@ -13,21 +13,36 @@
 
 # Third Party
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 
 def interpolate_models(
     model_path: str,
     trained_model_path: str,
     trained_model_weight: float = 0.5,
     output_model_path: str | None = None,
-    torch_dtype: str | None = "bfloat16",
+    torch_dtype: str | torch.dtype | None = "bfloat16",
 ) -> str:
     if output_model_path is None:
         output_model_path = f"{trained_model_path}_interp"
 
-    model_kwargs: dict[str, any] = {}
-    if torch_dtype is not None and torch_dtype != "auto":
-        model_kwargs["torch_dtype"] = torch_dtype
+    model_kwargs = {}
+    if torch_dtype is not None:
+        if isinstance(torch_dtype, str):
+            _torch_dtype = torch_dtype.lower()
+            if _torch_dtype == "auto":
+                model_kwargs["torch_dtype"] = "auto"
+            else:
+                _map = {
+                    "bfloat16": torch.bfloat16, "bf16": torch.bfloat16,
+                    "float16": torch.float16, "fp16": torch.float16,
+                    "float32": torch.float32, "fp32": torch.float32,
+                }
+                if _torch_dtype not in _map:
+                    raise ValueError(f"Unsupported --torch-dtype: {torch_dtype}")
+                model_kwargs["torch_dtype"] = _map[_torch_dtype]
+        else:
+            model_kwargs["torch_dtype"] = torch_dtype
 
     # load original model
     model = AutoModelForCausalLM.from_pretrained(
@@ -66,31 +81,31 @@ def parse_arguments():
         "--model-path",
         type=str,
         required=True,
-        help="path to the original model",
+        help="Path to the original model",
     )
     parser.add_argument(
         "--trained-model-path",
         type=str,
         required=True,
-        help="path to the trained model",
+        help="Path to the trained model",
     )
     parser.add_argument(
         "--trained-model-weight",
         type=float,
         default=0.5,
-        help="weight for the trained model",
+        help="Weight for the trained model",
     )
     parser.add_argument(
         "--output-model-path",
         type=str,
         default=None,
-        help="path to the output model",
+        help="Path to the output model",
     )
     parser.add_argument(
         "--torch-dtype",
         type=str,
         default="bfloat16",
-        help="torch dtype",
+        help="Torch dtype",
     )
     args = parser.parse_args()
     return args