Improve ludwig feature dict

.github/workflows/pytest.yml

@@ -254,7 +254,7 @@ jobs:
           # remove torch and ray from the dependencies so we can add them depending on the matrix args for the job.
           cat requirements.txt | sed '/^torch[>=<\b]/d' | sed '/^torchtext/d' | sed '/^torchvision/d' | sed '/^torchaudio/d' > requirements-temp && mv requirements-temp requirements.txt
           cat requirements_distributed.txt | sed '/^ray[\[]/d'
-          pip install torch==2.0.0 torchtext torchvision torchaudio
+          pip install torch==2.1.0 torchtext torchvision torchaudio
           pip install ray==2.3.0
           pip install '.[test]'
           pip list

ludwig/explain/captum.py

@@ -216,13 +216,13 @@ def explain(self) -> ExplanationsResult:
                 feat_to_token_attributions_global[feat_name] = token_attributions_global
 
             self.global_explanation.add(
-                input_features.keys(), total_attribution_global, feat_to_token_attributions_global
+                list(input_features.keys()), total_attribution_global, feat_to_token_attributions_global
             )
 
             for i, (feature_attributions, explanation) in enumerate(zip(total_attribution, self.row_explanations)):
                 # Add the feature attributions to the explanation object for this row.
                 explanation.add(
-                    input_features.keys(),
+                    list(input_features.keys()),
                     feature_attributions,
                     {k: v[i] for k, v in feat_to_token_attributions.items()},
                 )
@@ -245,7 +245,7 @@ def explain(self) -> ExplanationsResult:
             }
             # Prepend the negative class to the list of label explanations.
             self.global_explanation.add(
-                input_features.keys(), negated_attributions, negated_token_attributions, prepend=True
+                list(input_features.keys()), negated_attributions, negated_token_attributions, prepend=True
             )
 
             for explanation in self.row_explanations:
@@ -257,7 +257,9 @@ def explain(self) -> ExplanationsResult:
                     if fa.token_attributions is not None
                 }
                 # Prepend the negative class to the list of label explanations.
-                explanation.add(input_features.keys(), negated_attributions, negated_token_attributions, prepend=True)
+                explanation.add(
+                    list(input_features.keys()), negated_attributions, negated_token_attributions, prepend=True
+                )
 
             # TODO(travis): for force plots, need something similar to SHAP E[X]
             expected_values.append(0.0)

ludwig/explain/captum_ray.py

@@ -115,13 +115,13 @@ def explain(self) -> ExplanationsResult:
                     feat_to_token_attributions_global[feat_name] = token_attributions_global
 
                 self.global_explanation.add(
-                    input_features.keys(), total_attribution_global, feat_to_token_attributions_global
+                    list(input_features.keys()), total_attribution_global, feat_to_token_attributions_global
                 )
 
                 for i, (feature_attributions, explanation) in enumerate(zip(total_attribution, self.row_explanations)):
                     # Add the feature attributions to the explanation object for this row.
                     explanation.add(
-                        input_features.keys(),
+                        list(input_features.keys()),
                         feature_attributions,
                         {k: v[i] for k, v in feat_to_token_attributions.items()},
                     )
@@ -140,7 +140,7 @@ def explain(self) -> ExplanationsResult:
             }
             # Prepend the negative class to the list of label explanations.
             self.global_explanation.add(
-                input_features.keys(), negated_attributions, negated_token_attributions, prepend=True
+                list(input_features.keys()), negated_attributions, negated_token_attributions, prepend=True
             )
 
             for explanation in self.row_explanations:
@@ -152,7 +152,9 @@ def explain(self) -> ExplanationsResult:
                     if fa.token_attributions is not None
                 }
                 # Prepend the negative class to the list of label explanations.
-                explanation.add(input_features.keys(), negated_attributions, negated_token_attributions, prepend=True)
+                explanation.add(
+                    list(input_features.keys()), negated_attributions, negated_token_attributions, prepend=True
+                )
 
             # TODO(travis): for force plots, need something similar to SHAP E[X]
             expected_values.append(0.0)

ludwig/explain/gbm.py

@@ -55,11 +55,11 @@ def explain(self) -> ExplanationsResult:
 
         expected_values = []
         for _ in range(self.vocab_size):
-            self.global_explanation.add(base_model.input_features.keys(), feat_imp)
+            self.global_explanation.add(list(base_model.input_features.keys()), feat_imp)
 
             for explanation in self.row_explanations:
                 # Add the feature attributions to the explanation object for this row.
-                explanation.add(base_model.input_features.keys(), feat_imp)
+                explanation.add(list(base_model.input_features.keys()), feat_imp)
 
             # TODO:
             expected_values.append(0.0)

ludwig/features/feature_utils.py

@@ -14,7 +14,8 @@
 # limitations under the License.
 # ==============================================================================
 import re
-from typing import Dict, List, Optional, Tuple, Union
+from collections.abc import MutableMapping
+from typing import Iterator, List, Optional, Union
 
 import numpy as np
 import torch
@@ -157,7 +158,7 @@ def get_name_from_module_dict_key(key: str, feature_name_suffix_length: int = FE
     return name[:-feature_name_suffix_length]
 
 
-class LudwigFeatureDict(torch.nn.Module):
+class LudwigFeatureDict(torch.nn.Module, MutableMapping):
     """Torch ModuleDict wrapper that permits keys with any name.
 
     Torch's ModuleDict implementation doesn't allow certain keys to be used if they conflict with existing class
@@ -174,39 +175,26 @@ class LudwigFeatureDict(torch.nn.Module):
     def __init__(self):
         super().__init__()
         self.module_dict = torch.nn.ModuleDict()
-        self.internal_key_to_original_name_map = {}
 
-    def get(self, key) -> torch.nn.Module:
+    def __getitem__(self, key: str) -> torch.nn.Module:
         return self.module_dict[get_module_dict_key_from_name(key)]
 
-    def set(self, key: str, module: torch.nn.Module) -> None:
+    def __setitem__(self, key: str, value: torch.nn.Module) -> None:
         module_dict_key_name = get_module_dict_key_from_name(key)
-        self.internal_key_to_original_name_map[module_dict_key_name] = key
-        self.module_dict[module_dict_key_name] = module
+        self.module_dict[module_dict_key_name] = value
 
-    def __len__(self) -> int:
-        return len(self.module_dict)
-
-    def __next__(self) -> None:
-        return next(iter(self))
+    def __delitem__(self, key: str) -> None:
+        del self.module_dict[get_module_dict_key_from_name(key)]
 
-    def __iter__(self) -> None:
-        return iter(self.keys())
+    def __iter__(self) -> Iterator[str]:
+        return (get_name_from_module_dict_key(key) for key in self.module_dict)
 
-    def keys(self) -> List[str]:
-        return [
-            get_name_from_module_dict_key(feature_name)
-            for feature_name in self.internal_key_to_original_name_map.keys()
-        ]
-
-    def values(self) -> List[torch.nn.Module]:
-        return [module for _, module in self.module_dict.items()]
+    def __len__(self) -> int:
+        return len(self.module_dict)
 
-    def items(self) -> List[Tuple[str, torch.nn.Module]]:
-        return [
-            (get_name_from_module_dict_key(feature_name), module) for feature_name, module in self.module_dict.items()
-        ]
+    def set(self, key: str, value: torch.nn.Module) -> None:
+        self[key] = value
 
-    def update(self, modules: Dict[str, torch.nn.Module]) -> None:
-        for feature_name, module in modules.items():
-            self.set(feature_name, module)
+    def __hash__(self) -> int:
+        """Static hash value, because the object is mutable, but needs to be hashable for pytorch."""
+        return 1

ludwig/models/ecd.py

@@ -143,7 +143,7 @@ def forward(
         else:
             targets = None
 
-        assert list(inputs.keys()) == self.input_features.keys()
+        assert list(inputs.keys()) == list(self.input_features.keys())
 
         encoder_outputs = self.encode(inputs)
         combiner_outputs = self.combine(encoder_outputs)

ludwig/models/gbm.py

@@ -101,14 +101,14 @@ def forward(
     ) -> Dict[str, torch.Tensor]:
         # Invoke output features.
         output_logits = {}
-        output_feature_name = self.output_features.keys()[0]
+        output_feature_name = next(iter(self.output_features.keys()))
         output_feature = self.output_features.get(output_feature_name)
 
         # If `inputs` is a tuple, it should contain `(inputs, targets)`.
         if isinstance(inputs, tuple):
             inputs, _ = inputs
 
-        assert list(inputs.keys()) == self.input_features.keys()
+        assert list(inputs.keys()) == list(self.input_features.keys())
 
         # If the model has not been compiled, predict using the LightGBM sklearn iterface. Otherwise, use torch with
         # the Hummingbird compiled model. Notably, when compiling the model to torchscript, compiling with Hummingbird

ludwig/models/llm.py

@@ -65,13 +65,13 @@ def __iter__(self) -> None:
         return iter(self.obj.keys())
 
     def keys(self) -> List[str]:
-        return self.obj.keys()
+        return list(self.obj.keys())
 
     def values(self) -> List[torch.nn.Module]:
-        return self.obj.values()
+        return list(self.obj.values())
 
     def items(self) -> List[Tuple[str, torch.nn.Module]]:
-        return self.obj.items()
+        return list(self.obj.items())
 
     def update(self, modules: Dict[str, torch.nn.Module]) -> None:
         self.obj.update(modules)
@@ -148,7 +148,8 @@ def __init__(
         )
 
         # Extract the decoder object for the forward pass
-        self._output_feature_decoder = ModuleWrapper(self.output_features.items()[0][1])
+        decoder = next(iter(self.output_features.values()))
+        self._output_feature_decoder = ModuleWrapper(decoder)
 
         self.attention_masks = None
 
@@ -401,7 +402,7 @@ def _unpack_inputs(
         else:
             targets = None
 
-        assert list(inputs.keys()) == self.input_features.keys()
+        assert list(inputs.keys()) == list(self.input_features.keys())
 
         input_ids = self.get_input_ids(inputs)
         target_ids = self.get_target_ids(targets) if targets else None

ludwig/trainers/trainer_lightgbm.py

@@ -831,8 +831,8 @@ def _construct_lgb_datasets(
         validation_set: Optional["Dataset"] = None,  # noqa: F821
         test_set: Optional["Dataset"] = None,  # noqa: F821
     ) -> Tuple[lgb.Dataset, List[lgb.Dataset], List[str]]:
-        X_train = training_set.to_scalar_df(self.model.input_features.values())
-        y_train = training_set.to_scalar_df(self.model.output_features.values())
+        X_train = training_set.to_scalar_df(list(self.model.input_features.values()))
+        y_train = training_set.to_scalar_df(list(self.model.output_features.values()))
 
         # create dataset for lightgbm
         # keep raw data for continued training https://github.com/microsoft/LightGBM/issues/4965#issuecomment-1019344293
@@ -850,8 +850,8 @@ def _construct_lgb_datasets(
         eval_sets = [lgb_train]
         eval_names = [LightGBMTrainer.TRAIN_KEY]
         if validation_set is not None:
-            X_val = validation_set.to_scalar_df(self.model.input_features.values())
-            y_val = validation_set.to_scalar_df(self.model.output_features.values())
+            X_val = validation_set.to_scalar_df(list(self.model.input_features.values()))
+            y_val = validation_set.to_scalar_df(list(self.model.output_features.values()))
             try:
                 lgb_val = lgb.Dataset(X_val, label=y_val, reference=lgb_train, free_raw_data=False).construct()
             except lgb.basic.LightGBMError as e:
@@ -869,8 +869,8 @@ def _construct_lgb_datasets(
             pass
 
         if test_set is not None:
-            X_test = test_set.to_scalar_df(self.model.input_features.values())
-            y_test = test_set.to_scalar_df(self.model.output_features.values())
+            X_test = test_set.to_scalar_df(list(self.model.input_features.values()))
+            y_test = test_set.to_scalar_df(list(self.model.output_features.values()))
             try:
                 lgb_test = lgb.Dataset(X_test, label=y_test, reference=lgb_train, free_raw_data=False).construct()
             except lgb.basic.LightGBMError as e: