additional fixes

.github/workflows/push.yml

@@ -52,6 +52,7 @@ jobs:
           -config eole/tests/data/data.yaml \
           -save_data /tmp/eole.train.check \
           -n_sample 30 \
+          -model {"architecture": "rnn"} \
           -training '{"num_workers": 0, "bucket_size": 1024}' \
           -src_vocab /tmp/eole.vocab.src \
           -tgt_vocab /tmp/eole.vocab.tgt \
@@ -337,7 +338,7 @@ jobs:
           -tgt_vocab /tmp/eole.vocab.tgt \
           -src_vocab_size 1000 \
           -tgt_vocab_size 1000 \
-          -model '{"hidden_size": 10, "embeddings": {"word_vec_size": 5}}' \
+          -model '{"architecture": "rnn", "hidden_size": 10, "embeddings": {"word_vec_size": 5}}' \
           -training '{"batch_size": 10, "num_workers": 0, "bucket_size": 1024, "train_steps": 10, "model_path": "/tmp/eole.model", "save_checkpoint_steps": 10}' \
           -report_every 5
         sed -i '1s/^/new_tok\t100000000\n/' /tmp/eole.vocab.src
@@ -347,7 +348,6 @@ jobs:
           -tgt_vocab /tmp/eole.vocab.tgt \
           -src_vocab_size 1000 \
           -tgt_vocab_size 1000 \
-          -model '{"hidden_size": 10, "embeddings": {"word_vec_size": 5}}' \
           -training '{"batch_size": 10, "num_workers": 0, "bucket_size": 1024, "train_steps": 20, "train_from": "/tmp/eole.model/step_10", "save_checkpoint_steps": 10, "update_vocab": True, "reset_optim": "states"}' \
           -report_every 5
     - name: Test checkpoint vocabulary update with LM
@@ -366,7 +366,6 @@ jobs:
           -config eole/tests/data/lm_data.yaml \
           -src_vocab /tmp/eole.vocab.src \
           -tgt_vocab /tmp/eole.vocab.src \
-          -model '{"layers": 2, "hidden_size": 16, "transformer_ff": 64, "embeddings": {"word_vec_size": 16}, "encoder": None, "decoder": {"decoder_type": "transformer_lm", "heads": 4}}' \
           -training '{"batch_size": 10, "num_workers": 0, "bucket_size": 1024, "train_steps": 20, "train_from": "/tmp/lm.eole.model/step_10", "save_checkpoint_steps": 10, "update_vocab": True, "reset_optim": "states"}' \
           -src_vocab_size 1000 \
           -tgt_vocab_size 1000 \

eole/config/common.py

@@ -1,7 +1,6 @@
 # import torch
 from typing import List, Literal
-from pydantic import Field, computed_field, model_validator
-from functools import cached_property
+from pydantic import Field, model_validator
 from eole.config.config import Config
 
 # from eole.utils.logging import logger
@@ -33,8 +32,7 @@ class DistributedConfig(Config):
         default=60, description="Timeout for one GPU to wait for the others."
     )
 
-    @computed_field
-    @cached_property
+    @property
     def parallel_gpu(self) -> int:  # converted to a `property` by `computed_field`
         return self.world_size if self.parallel_mode == "tensor_parallel" else 1
 

eole/config/run.py

@@ -37,8 +37,13 @@ def get_model_path(self):
         return self.training.get_model_path()
 
     @classmethod
-    def get_defaults(cls):
-        return cls(src_vocab="dummy", tgt_vocab="dummy", data={}).model_dump()
+    def get_defaults(cls, architecture):
+        return cls(
+            src_vocab="dummy",
+            tgt_vocab="dummy",
+            data={},
+            model={"architecture": architecture},
+        ).model_dump()
 
     @field_validator("model", "training", mode="before")
     @classmethod
@@ -53,15 +58,11 @@ def str_to_dict(cls, v) -> Dict:
     @model_validator(mode="before")
     @classmethod
     def default_architecture(cls, data: Any) -> Any:
-        # somewhat dirty patch to make test_models pass, might do better later
-        # this explicit call to the field_validator should not be necessary
+        # only enforce default "custom" if some model settings are passed
         if "model" in data.keys():
             data["model"] = cls.str_to_dict(data["model"])
-        if isinstance(data.get("model", {}), dict):
-            if data.get("model", {}).get("architecture", None) is None:
-                if "model" not in data.keys():
-                    data["model"] = {"architecture": "custom"}
-                else:
+            if isinstance(data.get("model", {}), dict):
+                if data.get("model", {}).get("architecture", None) is None:
                     data["model"]["architecture"] = "custom"
         return data
 
@@ -91,7 +92,7 @@ class PredictConfig(
     src_subword_vocab: str | None = (
         None  # patch for CT2 inference engine (to improve later)
     )
-    model: ModelConfig
+    model: ModelConfig | None = None
 
     @model_validator(mode="after")
     def _validate_predict_config(self):

eole/decoders/ensemble.py

@@ -10,7 +10,7 @@
 import copy
 from eole.encoders.encoder import EncoderBase
 from eole.decoders.decoder import DecoderBase
-from eole.models import EncoderDecoderModel, get_model_class
+from eole.models import EncoderDecoderModel, BaseModel
 
 
 class EnsembleDecoderOutput(object):
@@ -191,7 +191,7 @@ def load_test_model(config, device_id=0):
     for i, model_path in enumerate(config.model_path):
         config2.model_path = [config.model_path[i]]
         print(config2.model)
-        vocabs, model, model_config = get_model_class(config2.model).load_test_model(
+        vocabs, model, model_config = BaseModel.load_test_model(
             config2, device_id, model_path=model_path
         )
         if shared_vocabs is None:

eole/predict/__init__.py

@@ -9,7 +9,7 @@
 from eole.predict.greedy_search import GreedySearch, GreedySearchLM
 from eole.predict.penalties import PenaltyBuilder
 from eole.decoders.ensemble import load_test_model as ensemble_load_test_model
-from eole.models import get_model_class
+from eole.models import BaseModel
 import codecs
 
 
@@ -44,7 +44,7 @@ def build_predictor(config, device_id=0, report_score=True, logger=None, out_fil
     load_test_model = (
         ensemble_load_test_model
         if len(config.model_path) > 1
-        else get_model_class(config.model).load_test_model
+        else BaseModel.load_test_model
     )
 
     vocabs, model, model_config = load_test_model(config, device_id)

eole/tests/pull_request_check.sh

@@ -89,6 +89,7 @@ rm -f -r $TMP_OUT_DIR/sample
 echo -n "[+] Testing NMT vocab? /transforms prepare..."
 ${PYTHON} eole/bin/main.py train \
             -config ${DATA_DIR}/data.yaml \
+            -model '{"architecture": "rnn"}' \
             -save_data $TMP_OUT_DIR/eole.train.check \
             -n_sample 30 \
             -overwrite \
@@ -271,7 +272,7 @@ ${PYTHON} eole/bin/main.py train \
             -src_vocab $TMP_OUT_DIR/eole.vocab.src \
             -tgt_vocab $TMP_OUT_DIR/eole.vocab.tgt \
             -src_vocab_size 1000 -tgt_vocab_size 1000 \
-            -model '{"hidden_size": 10, "embeddings": {"word_vec_size": 5}}' \
+            -model '{"architecture": "rnn", "hidden_size": 10, "embeddings": {"word_vec_size": 5}}' \
             -training '{"batch_size": 10, "num_workers": 0, "bucket_size": 1024, "train_steps": 10, "model_path": "'"$TMP_OUT_DIR"'/eole.model", "save_checkpoint_steps": 10}' \
             -report_every 5 \
             >> ${LOG_FILE} 2>&1
@@ -281,7 +282,6 @@ ${PYTHON} eole/bin/main.py train \
             -src_vocab $TMP_OUT_DIR/eole.vocab.src \
             -tgt_vocab $TMP_OUT_DIR/eole.vocab.tgt \
             -src_vocab_size 1000 -tgt_vocab_size 1000 \
-            -model '{"hidden_size": 10, "embeddings": {"word_vec_size": 5}}' \
             -training '{"batch_size": 10, "num_workers": 0, "bucket_size": 1024, "train_steps": 20, "train_from": "'"$TMP_OUT_DIR"'/eole.model/step_10", "save_checkpoint_steps": 10, "update_vocab": True, "reset_optim": "states"}' \
             -report_every 5 \
             >> ${LOG_FILE} 2>&1

eole/tests/test_models.py

@@ -11,6 +11,7 @@
 from eole.utils.misc import sequence_mask
 from eole.config.run import TrainConfig
 from eole.config.data import Dataset
+from eole.config.models import CustomModelConfig
 
 # In theory we should probably call ModelConfig here,
 # but we can't because model building relies on some params
@@ -22,6 +23,7 @@
     },  # actual file path (tested in validate)
     src_vocab="dummy",
     share_vocab=True,
+    model=CustomModelConfig(),
 )  # now required by validation
 
 
@@ -113,7 +115,7 @@ def encoder_forward(self, opt, source_l=3, bsize=1):
 
         # Initialize vectors to compare size with
         test_hid = torch.zeros(
-            self.opt.model.encoder.layers, bsize, opt.model.encoder.hidden_size
+            opt.model.encoder.layers, bsize, opt.model.encoder.hidden_size
         )
         test_out = torch.zeros(bsize, source_l, opt.model.decoder.hidden_size)
 
@@ -192,17 +194,26 @@ def test_method(self):
 """
 # opt.brnn = False # deprecated and not used here
 
-test_embeddings = [[("model", {"architecture": "transformer"})]]
+test_embeddings = [
+    [("model", {"architecture": "rnn"})],
+    [("model", {"architecture": "transformer"})],
+]
 
 for p in test_embeddings:
     _add_test(p, "embeddings_forward")
 
 tests_encoder = [
-    # not supported anymore
-    # [
-    #     # ("encoder_type", "mean"),
-    #     ("model", {"architecture": "custom", "encoder": {"encoder_type": "mean"}})
-    # ],
+    [
+        # ("encoder_type", "mean"),
+        (
+            "model",
+            {
+                "architecture": "custom",
+                "encoder": {"encoder_type": "mean"},
+                "decoder": {"decoder_type": "rnn"},
+            },
+        )
+    ],
     # [('encoder_type', 'transformer'),
     # ('word_vec_size', 16), ('hidden_size', 16)],
 ]

eole/tests/test_transform.py

@@ -14,6 +14,7 @@
 from eole.transforms.bart import BARTNoising
 from eole.config.run import TrainConfig, PredictConfig
 from eole.config.data import Dataset
+from eole.config.models import CustomModelConfig
 
 
 class TestTransform(unittest.TestCase):
@@ -43,6 +44,7 @@ def test_vocab_required_transform(self):
             share_vocab=True,
             seed=-1,
             transforms_configs={"switchout": {"switchout_temperature": 1.0}},
+            model=CustomModelConfig(),
         )
         # transforms that require vocab will not create if not provide vocab
         transforms = make_transforms(opt, transforms_cls, vocabs=None)
@@ -67,7 +69,9 @@ def test_transform_specials(self):
         """
         )
         # opt = Namespace(data=corpora)
-        opt = TrainConfig(data=corpora, src_vocab="dummy", share_vocab=True)
+        opt = TrainConfig(
+            data=corpora, src_vocab="dummy", share_vocab=True, model=CustomModelConfig()
+        )
         specials = get_specials(opt, transforms_cls)
         specials_expected = {"src": ["｟_pf_src｠"], "tgt": ["｟_pf_tgt｠"]}
         self.assertEqual(specials, specials_expected)
@@ -87,7 +91,13 @@ def test_transform_pipe(self):
         """
         )
         # opt = Namespace(data=corpora, seed=-1)
-        opt = TrainConfig(data=corpora, seed=-1, src_vocab="dummy", share_vocab=True)
+        opt = TrainConfig(
+            data=corpora,
+            seed=-1,
+            src_vocab="dummy",
+            share_vocab=True,
+            model=CustomModelConfig(),
+        )
         prefix_transform = prefix_cls(opt)
         prefix_transform.warm_up()
         # 2. Init second transform in the pipe
@@ -102,6 +112,7 @@ def test_transform_pipe(self):
             transforms_configs={
                 "filtertoolong": {"src_seq_length": 4, "tgt_seq_length": 4}
             },
+            model=CustomModelConfig(),
         )
         filter_transform = filter_cls(opt)
         # 3. Sequential combine them into a transform pipe
@@ -137,7 +148,13 @@ def test_prefix(self):
         """
         )
         # opt = Namespace(data=corpora, seed=-1)
-        opt = TrainConfig(data=corpora, seed=-1, src_vocab="dummy", share_vocab=True)
+        opt = TrainConfig(
+            data=corpora,
+            seed=-1,
+            src_vocab="dummy",
+            share_vocab=True,
+            model=CustomModelConfig(),
+        )
         prefix_transform = prefix_cls(opt)
         prefix_transform.warm_up()
         self.assertIn("trainset", prefix_transform.prefix_dict)
@@ -498,6 +515,7 @@ def test_tokendrop(self):
             share_vocab=True,
             seed=3434,
             transforms_configs={"tokendrop": {"tokendrop_temperature": 0.1}},
+            model=CustomModelConfig(),
         )
         tokendrop_transform = tokendrop_cls(opt)
         tokendrop_transform.warm_up()
@@ -523,6 +541,7 @@ def test_tokenmask(self):
             share_vocab=True,
             seed=3434,
             transforms_configs={"tokenmask": {"tokenmask_temperature": 0.1}},
+            model=CustomModelConfig(),
         )
         tokenmask_transform = tokenmask_cls(opt)
         tokenmask_transform.warm_up()
@@ -548,6 +567,7 @@ def test_switchout(self):
             share_vocab=True,
             seed=3434,
             transforms_configs={"switchout": {"switchout_temperature": 0.1}},
+            model=CustomModelConfig(),
         )
         switchout_transform = switchout_cls(opt)
         with self.assertRaises(ValueError):

eole/train_single.py

@@ -109,7 +109,9 @@ def configure_process(config, device_id):
 
 def update_config_with_checkpoint(config, checkpoint=None):
     if checkpoint is not None:
-        defaults = TrainConfig.get_defaults()
+        defaults = TrainConfig.get_defaults(
+            architecture=checkpoint["config"].model.architecture
+        )
         checkpoint_non_defaults = recursive_model_fields_set(checkpoint["config"])
         new_config = recursive_model_fields_set(config)
         # override data explicitly,