updated the data type of meds_to_cehrbert_conversion_type in the dataclass field definition (#112)

ChaoPang · web-flow · commit ee58b106debd · 2025-05-19T14:04:56.000-04:00
diff --git a/src/cehrbert/runners/hf_runner_argument_dataclass.py b/src/cehrbert/runners/hf_runner_argument_dataclass.py
@@ -117,10 +117,10 @@ class DataTrainingArguments:
     # TODO: Python 3.9/10 do not support dynamic unpacking, we have to manually provide the entire
     #  list right now.
     meds_to_cehrbert_conversion_type: Literal[
-        MedsToCehrBertConversionType[MedsToBertMimic4.__name__],
-        MedsToCehrBertConversionType[MedsToCehrbertOMOP.__name__],
+        MedsToBertMimic4.__name__,
+        MedsToCehrbertOMOP.__name__,
     ] = dataclasses.field(
-        default=MedsToCehrBertConversionType[MedsToBertMimic4.__name__],
+        default=MedsToCehrbertOMOP.__name__,
         metadata={
             "help": "The MEDS to CEHRBERT conversion type e.g. MedsToBertMimic4",
             "choices": f"choices={[e for e in MedsToCehrBertConversionType.__members__]}",
diff --git a/tests/integration_tests/runners/hf_cehrbert_pretrain_runner_meds_test.py b/tests/integration_tests/runners/hf_cehrbert_pretrain_runner_meds_test.py
@@ -58,7 +58,7 @@ def setUp(self):
             "--meds_exclude_tables",
             "condition",
             "--meds_to_cehrbert_conversion_type",
-            "MedsToCehrBertConversionType.MedsToCehrbertOMOP",
+            "MedsToCehrbertOMOP",
         ]
 
     def tearDown(self):

Original file line number	Diff line number	Diff line change
`@@ -58,7 +58,7 @@ def setUp(self):`
`58`	`58`	`"--meds_exclude_tables",`
`59`	`59`	`"condition",`
`60`	`60`	`"--meds_to_cehrbert_conversion_type",`
`61`		`- "MedsToCehrBertConversionType.MedsToCehrbertOMOP",`
	`61`	`+ "MedsToCehrbertOMOP",`
`62`	`62`	`]`
`63`	`63`
`64`	`64`	`def tearDown(self):`