remove model.model in ckpt loading

torchchat/cli/builder.py

@@ -388,11 +388,9 @@ def _load_model_default(builder_args, only_config=False):
             builder_args.device
         ):
             model = Model.from_params(builder_args.params_path)
-        state_dict = flamingo_meta_to_tune(checkpoint)
-        model.model.load_state_dict(state_dict)
-    else:
-        checkpoint = {"model." + k: v for k, v in checkpoint.items()}
-        model.load_state_dict(checkpoint, assign=True, strict=True)
+        checkpoint = flamingo_meta_to_tune(checkpoint)
+
+    model.load_state_dict(checkpoint, assign=True, strict=True)
 
     return model
 

torchchat/model.py

@@ -442,6 +442,8 @@ def __init__(self, config: ModelArgs) -> None:
         # It should be assigned in the actual model implementation, if any.
         self.text_transformer_args = None
 
+        self._register_load_state_dict_pre_hook(self._load_model_state_dict)
+
     def build_model(self) -> nn.Module:
         """
         Builds a model based on the provided configuration.
@@ -468,6 +470,13 @@ def _replace_known_params(self, params):
                 params[key] = patterns[value]
         return params
 
+    def _load_model_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        # 修改 state dict 中的键值
+        for key in list(state_dict.keys()):
+            new_key = 'model.' + key
+            state_dict[new_key] = state_dict.pop(key)
+        return state_dict
+
     @abstractmethod
     def forward(self, *args, **kwargs):
         raise NotImplementedError("forward method is not implemented")

torchchat/utils/gguf_loader.py

@@ -47,7 +47,6 @@ def _convert_gguf_tensor_name_to_llama_nn(gguf_name: str) -> str:
     result = copy.deepcopy(gguf_name)
     for gguf_string, replacement in _name_replacements:
         result = result.replace(gguf_string, replacement)
-    result = "model." + result
     return result