chg: Implemented pretrained weight loading in get_model + tests

Author: joshhan619

ltsm/common/base_training_pipeline.py

@@ -52,6 +52,7 @@ class TrainingConfig:
         "data": "ETTh1",
         "features": "MS",
         "prompt_data_path": "./weather.csv",
+        "hf_hub_model": None
     }
 
     def __init__(self, model_config: PretrainedConfig, **kwargs):
@@ -126,7 +127,7 @@ def __init__(self,
         self.config = config
 
         if not model:
-            self.model = get_model(config.model_config, config.train_params["model"], config.train_params["local_pretrain"])
+            self.model = get_model(config.model_config, config.train_params["model"], config.train_params["local_pretrain"], config.train_params["hf_hub_model"])
 
         if self.config.train_params["lora"]:
             peft_config = LoraConfig(

ltsm/models/__init__.py

@@ -29,14 +29,15 @@ def register_model(module, module_name: str):
 register_model(DLinear, 'DLinear')
 register_model(Informer, 'Informer')
 
-def get_model(config: PretrainedConfig, model_name: str, local_pretrain: str = None) -> PreTrainedModel:
+def get_model(config: PretrainedConfig, model_name: str, local_pretrain: str = None, hf_hub_model: str = None) -> PreTrainedModel:
     """
     Factory method to create a model by name.
     
     Args:
         config (PreTrainedConfig): The configuration for the model.
         model_name (str): The name of the model to instantiate.
         local_pretrain (bool): If True, load the model from a local pretraining path.
+        hf_hub_model (str): The Hugging Face Hub model name.
     
     Returns:
         torch.nn.Module: Instantiated model.
@@ -47,6 +48,10 @@ def get_model(config: PretrainedConfig, model_name: str, local_pretrain: str = N
     if model_name not in model_dict:
         raise ValueError(f"Model {model_name} is not registered. Available models: {list(model_dict.keys())}")
 
+    # Load pretrained weights if hf_hub_model is provided
+    if hf_hub_model is not None:
+        return model_dict[model_name].from_pretrained(hf_hub_model, config)
+
     # Check for local pretraining
     if local_pretrain is None or local_pretrain == "None":
         return model_dict[model_name](config)

tests/model/DLinear_test.py tests/models/DLinear_test.pytests/model/DLinear_test.py renamed to tests/models/DLinear_test.py

@@ -0,0 +1,44 @@
+import pytest
+from transformers import PretrainedConfig, PreTrainedModel
+from ltsm.models import register_model, get_model, model_dict
+
+def test_register_model(mocker):
+    mock_model = mocker.MagicMock(spec=PreTrainedModel)
+    register_model(mock_model, "MockModel1")
+    assert "MockModel1" in model_dict
+    assert model_dict["MockModel1"] == mock_model
+
+    with pytest.raises(AssertionError, match="Reader MockModel1 already registered"):
+        register_model(mock_model, "MockModel1")
+
+def test_get_model(mocker):
+    mock_model = mocker.MagicMock(spec=PreTrainedModel)
+    mock_config = mocker.MagicMock(spec=PretrainedConfig)
+    register_model(mock_model, "MockModel2")
+    
+    instance = get_model(mock_config, "MockModel2")
+    mock_model.assert_called_once_with(mock_config)
+    assert isinstance(instance, mocker.MagicMock)
+
+def test_get_model_invalid_name():
+    with pytest.raises(ValueError, match="Model NonExistentModel is not registered"):
+        get_model(PretrainedConfig(), "NonExistentModel")
+
+def test_get_model_local_pretrain(mocker):
+    mock_from_pretrained = mocker.patch("transformers.PretrainedConfig.from_pretrained")
+    mock_model = mocker.MagicMock(spec=PreTrainedModel)
+    register_model(mock_model, "MockModel3")
+    
+    mock_from_pretrained.return_value = mocker.MagicMock()
+    instance = get_model(PretrainedConfig(), "MockModel3", local_pretrain="path/to/pretrained")
+    mock_model.from_pretrained.assert_called_once_with("path/to/pretrained", mock_from_pretrained.return_value)
+    assert isinstance(instance, mocker.MagicMock)
+
+def test_get_model_hf_hub(mocker):
+    mock_from_pretrained = mocker.patch("transformers.PreTrainedModel.from_pretrained")
+    mock_model = mocker.MagicMock(spec=PreTrainedModel)
+    register_model(mock_model, "MockModel4")
+    
+    instance = get_model(PretrainedConfig(), "MockModel4", hf_hub_model="mock-hub-model")
+    mock_model.from_pretrained.assert_called_once_with("mock-hub-model", PretrainedConfig())
+    assert isinstance(instance, mocker.MagicMock)