optimize sampling utils

data/text_sampling/text_sampling.py

@@ -664,7 +664,6 @@ def get_sample_dict(self, sample: pd.Series, template: str):
             multiple_choice_enum_idx = multiple_choice_enum_idx[0]  # unpack list
             multiple_choice_enum = input_variables[multiple_choice_enum_idx]
 
-
             # get multiple_choice_var
             multiple_choice_var_idx = [
                 i for i, x in enumerate(input_variables) if x.endswith("%")

pyproject.toml

@@ -34,7 +34,7 @@ dataset_creation = [
 chemnlp-generate-meta = "chemnlp.data.meta_yaml_generator:cli"
 chemnlp-augment-meta = "chemnlp.data.meta_yaml_augmenter:cli"
 chemnlp-sample = "chemnlp.data.sampler_cli:cli"
-chemlp-add-random-split-column = "chemnlp.data.utils:add_random_split_column_cli"
+chemnlp-add-random-split-column = "chemnlp.data.utils:add_random_split_column_cli"
 
 [tool.setuptools_scm]
 version_scheme = "post-release"

src/chemnlp/data/meta_yaml_augmentor.py

@@ -152,8 +152,10 @@ def _cli(data_dir: str, model: str = "gpt-4o", override: bool = False):
 
     return augmented_meta_yaml
 
+
 def cli():
     fire.Fire(_cli)
 
+
 if __name__ == "__main__":
     fire.Fire(_cli)

src/chemnlp/data/meta_yaml_generator.py

@@ -198,6 +198,7 @@ def _cli(
 def cli():
     fire.Fire(_cli)
 
+
 # Example usage
 if __name__ == "__main__":
     fire.Fire(_cli)

src/chemnlp/data/sampler.py

@@ -187,7 +187,9 @@ def _wrap_identifier(self, identifier: str, value: str) -> str:
         except ValueError:
             identifier_type = None
 
-        if identifier_type and identifier_type not in self.config.get('excluded_from_wrapping', []):
+        if identifier_type and identifier_type not in self.config.get(
+            "excluded_from_wrapping", []
+        ):
             return f"[BEGIN_{identifier_type}]{value}[END_{identifier_type}]"
         return value
 

src/chemnlp/data/sampler_cli.py

@@ -93,7 +93,7 @@ def process_dataset(
         "multiple_choice_benchmarking_format": None,
         "wrap_identifiers": wrap_identifiers,
         "benchmarking_templates": benchmarking,
-        "excluded_from_wrapping": ['Other']
+        "excluded_from_wrapping": ["Other"],
     }
 
     templates = meta["templates"]
@@ -172,8 +172,10 @@ def main(
         wrap_identifiers,
     )
 
+
 def cli():
     fire.Fire(main)
 
+
 if __name__ == "__main__":
     fire.Fire(main)

src/chemnlp/data/utils.py

@@ -26,12 +26,14 @@ def add_random_split_column(df):
 
     return df
 
+
 def _add_random_split_column(file):
     df = pd.read_csv(file)
     df = add_random_split_column(df)
     df.to_csv(file, index=False)
 
-def add_random_split_column_cli(file: str):
+
+def add_random_split_column_cli():
     fire.Fire(_add_random_split_column)
 
 

tests/data/test_sampler.py

@@ -111,7 +111,7 @@ def sample_config():
         "multiple_choice_rnd_symbols": ["", ".)", ")"],
         "multiple_choice_benchmarking_templates": False,
         "multiple_choice_benchmarking_format": None,
-                "excluded_from_wrapping": ['Other']
+        "excluded_from_wrapping": ["Other"],
     }
 
 
@@ -123,7 +123,7 @@ def sample_config_with_wrapping():
         "multiple_choice_benchmarking_templates": False,
         "multiple_choice_benchmarking_format": None,
         "wrap_identifiers": True,
-                "excluded_from_wrapping": ['Other']
+        "excluded_from_wrapping": ["Other"],
     }