config bugfix

jlunder00 · Nov 25, 2024 · 82c65a0 · 82c65a0
1 parent 8ca0151
commit 82c65a0
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 2 deletions.
diff --git a/TMN_DataGen/dataset_generator.py b/TMN_DataGen/dataset_generator.py
@@ -88,6 +88,7 @@ def generate_dataset(
         output_path: str,
         parser_config: Optional[Union[str, Dict]] = None,
         preprocessing_config: Optional[Union[str, Dict]] = None,
+        feature_config: Optional[Union[str, Dict]] = None,
         verbosity: str = 'normal',
         override_pkg_config: Optional[Union[str, Dict]] = None,
         show_progress: bool = True
@@ -109,6 +110,7 @@ def generate_dataset(
         config, pkg_config = self._load_configs(
             parser_config,
             preprocessing_config,
+            feature_config,
             verbosity,
             override_pkg_config
         )

diff --git a/setup.py b/setup.py
@@ -10,7 +10,7 @@
 
 setup(
     name="TMN_DataGen",
-    version='0.6.0',
+    version='0.6.1',
     description="Tree Matching Network Data Generator",
     author="toast",
     packages=find_packages(),

diff --git a/tests/test_dataset_generator.py b/tests/test_dataset_generator.py
@@ -39,7 +39,11 @@ def test_config_override(sample_data, default_config, tmp_path):
         sentence_pairs=sample_data['sentence_pairs'],
         labels=sample_data['labels'], 
         output_path=str(output_path),
-        parser_config=config,
+        parser_config={'parser': config['parser']},
+        feature_config= {
+            'feature_extraction':config['feature_extraction'],
+            'feature_mapping':config['feature_mapping']
+        },
         verbosity='debug'
     )