cehrgpt 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

cehrgpt/data/hf_cehrgpt_dataset_mapping.py

@@ -1,8 +1,10 @@
 import datetime
-from typing import Any, Dict, Generator, Optional
+from collections import defaultdict
+from typing import Any, Dict, Generator, List, Optional, Union
 
 import numpy as np
 import pandas as pd
+from cehrbert.data_generators.hf_data_generator import UNKNOWN_VALUE
 from cehrbert.data_generators.hf_data_generator.hf_dataset_mapping import (
     ED_VISIT_TYPE_CODES,
     INPATIENT_VISIT_TYPE_CODES,
@@ -15,9 +17,17 @@ from cehrbert.data_generators.hf_data_generator.hf_dataset_mapping import (
 )
 from cehrbert.med_extension.schema_extension import Event
 from cehrbert.runners.hf_runner_argument_dataclass import DataTrainingArguments
+from cehrbert_data.const.artificial_tokens import (
+    DISCHARGE_UNKNOWN_TOKEN,
+    GENDER_UNKNOWN_TOKEN,
+    RACE_UNKNOWN_TOKEN,
+    VISIT_UNKNOWN_TOKEN,
+)
 from cehrbert_data.const.common import NA
 from cehrbert_data.decorators.patient_event_decorator_base import get_att_function
+from datasets.formatting.formatting import LazyBatch
 from dateutil.relativedelta import relativedelta
+from pandas import Series
 
 from cehrgpt.models.tokenization_hf_cehrgpt import (
     NONE_BIN,
@@ -25,6 +35,17 @@ from cehrgpt.models.tokenization_hf_cehrgpt import (
     CehrGptTokenizer,
 )
 
+CEHRGPT_COLUMNS = [
+    "concept_ids",
+    "concept_value_masks",
+    "number_as_values",
+    "concept_as_values",
+    "is_numeric_types",
+    "concept_values",
+    "units",
+    "epoch_times",
+]
+
 
 def convert_date_to_posix_time(index_date: datetime.date) -> float:
     return datetime.datetime.combine(
@@ -32,7 +53,36 @@ def convert_date_to_posix_time(index_date: datetime.date) -> float:
     ).timestamp()
 
 
-class MedToCehrGPTDatasetMapping(DatasetMapping):
+class DatasetMappingDecorator(DatasetMapping):
+
+    def batch_transform(
+        self, records: Union[LazyBatch, Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
+        """
+        Drop index_date if it contains None.
+
+        :param records:
+        :return:
+        """
+        if isinstance(records, LazyBatch):
+            table = records.pa_table
+
+            if "index_date" in table.column_names:
+                index_col = table.column("index_date")
+                if index_col.null_count > 0:
+                    table = table.drop(["index_date"])
+            records = LazyBatch(pa_table=table, formatter=records.formatter)
+        else:
+            if "index_date" in records:
+                if pd.isna(records["index_date"][0]):
+                    del records["index_date"]
+        return super().batch_transform(records=records)
+
+    def transform(self, record: Dict[str, Any]) -> Union[Dict[str, Any], Series]:
+        raise NotImplemented("Must be implemented")
+
+
+class MedToCehrGPTDatasetMapping(DatasetMappingDecorator):
     def __init__(
         self,
         data_args: DataTrainingArguments,
@@ -65,6 +115,7 @@ class MedToCehrGPTDatasetMapping(DatasetMapping):
     def _update_cehrgpt_record(
         cehrgpt_record: Dict[str, Any],
         code: str,
+        time: datetime.datetime,
         concept_value_mask: int = 0,
         number_as_value: float = 0.0,
         concept_as_value: str = "0",
@@ -77,6 +128,7 @@ class MedToCehrGPTDatasetMapping(DatasetMapping):
         cehrgpt_record["concept_as_values"].append(concept_as_value)
         cehrgpt_record["units"].append(unit)
         cehrgpt_record["is_numeric_types"].append(is_numeric_type)
+        cehrgpt_record["epoch_times"].append(time.timestamp())
 
     def transform(self, record: Dict[str, Any]) -> Dict[str, Any]:
         cehrgpt_record = {
@@ -87,13 +139,16 @@ class MedToCehrGPTDatasetMapping(DatasetMapping):
             "concept_as_values": [],
             "units": [],
             "is_numeric_types": [],
+            "epoch_times": [],
         }
         # Extract the demographic information
         birth_datetime = record["birth_datetime"]
         if isinstance(birth_datetime, pd.Timestamp):
             birth_datetime = birth_datetime.to_pydatetime()
         gender = record["gender"]
+        gender = GENDER_UNKNOWN_TOKEN if gender == UNKNOWN_VALUE else gender
         race = record["race"]
+        race = RACE_UNKNOWN_TOKEN if race == UNKNOWN_VALUE else race
         visits = record["visits"]
         # This indicates this is columnar format
         if isinstance(visits, dict):
@@ -108,10 +163,12 @@ class MedToCehrGPTDatasetMapping(DatasetMapping):
         )
         year_str = f"year:{str(first_visit_start_datetime.year)}"
         age_str = f"age:{str(relativedelta(first_visit_start_datetime, birth_datetime).years)}"
-        self._update_cehrgpt_record(cehrgpt_record, year_str)
-        self._update_cehrgpt_record(cehrgpt_record, age_str)
-        self._update_cehrgpt_record(cehrgpt_record, gender)
-        self._update_cehrgpt_record(cehrgpt_record, race)
+        self._update_cehrgpt_record(
+            cehrgpt_record, year_str, first_visit_start_datetime
+        )
+        self._update_cehrgpt_record(cehrgpt_record, age_str, first_visit_start_datetime)
+        self._update_cehrgpt_record(cehrgpt_record, gender, first_visit_start_datetime)
+        self._update_cehrgpt_record(cehrgpt_record, race, first_visit_start_datetime)
 
         # Use a data cursor to keep track of time
         datetime_cursor: Optional[datetime.datetime] = None
@@ -146,6 +203,7 @@ class MedToCehrGPTDatasetMapping(DatasetMapping):
                 self._update_cehrgpt_record(
                     cehrgpt_record,
                     code=self._time_token_function(time_delta),
+                    time=visit_start_datetime,
                 )
 
             datetime_cursor = visit_start_datetime
@@ -153,11 +211,13 @@ class MedToCehrGPTDatasetMapping(DatasetMapping):
             self._update_cehrgpt_record(
                 cehrgpt_record,
                 code="[VS]",
+                time=datetime_cursor,
             )
             # Add a visit type token
             self._update_cehrgpt_record(
                 cehrgpt_record,
                 code=visit_type,
+                time=datetime_cursor,
             )
             # We need to insert an inpatient hour token right after the visit type, we calculate the hour interval
             # with respect to the midnight of the day
@@ -167,6 +227,7 @@ class MedToCehrGPTDatasetMapping(DatasetMapping):
                     self._update_cehrgpt_record(
                         cehrgpt_record,
                         code=f"i-H{datetime_cursor.hour}",
+                        time=datetime_cursor,
                     )
 
             # Keep track of the existing outpatient events, we don't want to add them again
@@ -185,6 +246,10 @@ class MedToCehrGPTDatasetMapping(DatasetMapping):
                 concept_value_mask = int(
                     numeric_value is not None or text_value is not None
                 )
+                if numeric_value is None and text_value is not None:
+                    if text_value.isnumeric():
+                        numeric_value = float(text_value)
+
                 is_numeric_type = int(numeric_value is not None)
                 code = replace_escape_chars(e["code"])
 
@@ -208,6 +273,7 @@ class MedToCehrGPTDatasetMapping(DatasetMapping):
                         self._update_cehrgpt_record(
                             cehrgpt_record,
                             code=f"i-{self._inpatient_time_token_function(time_diff_days)}",
+                            time=event_time,
                         )
 
                     if self._include_inpatient_hour_token:
@@ -226,6 +292,7 @@ class MedToCehrGPTDatasetMapping(DatasetMapping):
                             self._update_cehrgpt_record(
                                 cehrgpt_record,
                                 code=f"i-H{time_diff_hours}",
+                                time=event_time,
                             )
 
                 if event_identity in existing_duplicate_events:
@@ -234,6 +301,7 @@ class MedToCehrGPTDatasetMapping(DatasetMapping):
                 self._update_cehrgpt_record(
                     cehrgpt_record,
                     code=code,
+                    time=event_time,
                     concept_value_mask=concept_value_mask,
                     unit=unit,
                     number_as_value=numeric_value if numeric_value else 0.0,
@@ -262,17 +330,24 @@ class MedToCehrGPTDatasetMapping(DatasetMapping):
                     # facility event
                     discharge_facility = get_value(visit, "discharge_facility")
                     if not discharge_facility:
-                        discharge_facility = "0"
-
+                        discharge_facility = DISCHARGE_UNKNOWN_TOKEN
+                    else:
+                        discharge_facility = (
+                            DISCHARGE_UNKNOWN_TOKEN
+                            if discharge_facility == UNKNOWN_VALUE
+                            else discharge_facility
+                        )
                     self._update_cehrgpt_record(
                         cehrgpt_record,
                         code=discharge_facility,
+                        time=datetime_cursor,
                     )
 
             # Reuse the age and date calculated for the last event in the patient timeline
             self._update_cehrgpt_record(
                 cehrgpt_record,
                 code="[VE]",
+                time=datetime_cursor,
             )
 
         # Generate the orders of the concepts that the cehrbert dataset mapping function expects
@@ -284,17 +359,21 @@ class MedToCehrGPTDatasetMapping(DatasetMapping):
         cehrgpt_record["num_of_concepts"] = len(cehrgpt_record["concept_ids"])
         cehrgpt_record["num_of_visits"] = len(visits)
 
-        if record.get("index_date", None):
+        if record.get("index_date", None) is not None:
             cehrgpt_record["index_date"] = record["index_date"]
-        if record.get("label", None):
+        if record.get("label", None) is not None:
             cehrgpt_record["label"] = record["label"]
-        if record.get("age_at_index", None):
+        if record.get("age_at_index", None) is not None:
             cehrgpt_record["age_at_index"] = record["age_at_index"]
 
+        assert len(cehrgpt_record["epoch_times"]) == len(
+            cehrgpt_record["concept_ids"]
+        ), "The number of time stamps must match with the number of concepts in the sequence"
+
         return cehrgpt_record
 
 
-class HFCehrGptTokenizationMapping(DatasetMapping):
+class HFCehrGptTokenizationMapping(DatasetMappingDecorator):
     def __init__(
         self,
         concept_tokenizer: CehrGptTokenizer,
@@ -308,9 +387,46 @@ class HFCehrGptTokenizationMapping(DatasetMapping):
             "is_numeric_types",
         ]
 
+    def filter_out_invalid_tokens(self, record: Dict[str, Any]) -> Dict[str, Any]:
+        column_names = []
+        seq_length = len(record["concept_ids"])
+
+        # We can't have "0" as a token in the tokenizer because it would break tokenization for "Race/0", "Visit/0"
+        # This is a pre-caution
+        if "0" in record["concept_ids"]:
+            if isinstance(record["concept_ids"], np.ndarray):
+                record["concept_ids"][record["concept_ids"] == "0"] = "Unknown"
+            else:
+                record["concept_ids"] = [
+                    "Unknown" if x == "0" else x for x in record["concept_ids"]
+                ]
+
+        for k, v in record.items():
+            if k not in CEHRGPT_COLUMNS:
+                continue
+            if isinstance(v, (list, np.ndarray)) and len(v) == seq_length:
+                column_names.append(k)
+        valid_concept_ids = self._concept_tokenizer.get_vocab().keys()
+        valid_indices = [
+            idx
+            for idx, concept_id in enumerate(record["concept_ids"])
+            if concept_id in valid_concept_ids
+        ]
+        if len(valid_indices) != len(record["concept_ids"]):
+            for column in column_names:
+                values = record[column]
+                record[column] = [values[idx] for idx in valid_indices]
+        return record
+
     def transform(self, record: Dict[str, Any]) -> Dict[str, Any]:
+        # Remove the tokens from patient sequences that do not exist in the tokenizer
+        record = self.filter_out_invalid_tokens(record)
         # If any concept has a value associated with it, we normalize the value
         record["input_ids"] = self._concept_tokenizer.encode(record["concept_ids"])
+        assert len(record["input_ids"]) == len(record["concept_ids"]), (
+            "The number of tokens must equal to the number of concepts\n"
+            f"decoded concept_ids: {self._concept_tokenizer.decode(record['input_ids'], skip_special_tokens=False)}"
+        )
         record["value_indicators"] = record["concept_value_masks"]
         if "number_as_values" not in record or "concept_as_values" not in record:
             record["number_as_values"] = [
@@ -391,3 +507,89 @@ class HFFineTuningMapping(HFCehrGptTokenizationMapping):
         columns = super().remove_columns()
         columns.append("label")
         return columns
+
+
+class ExtractTokenizedSequenceDataMapping:
+    def __init__(
+        self,
+        person_index_date_map: Dict[int, List[Dict[str, Any]]],
+        observation_window: int = 0,
+    ):
+        self.person_index_date_map = person_index_date_map
+        self.observation_window = observation_window
+
+    def _calculate_prediction_start_time(self, prediction_time: float):
+        if self.observation_window and self.observation_window > 0:
+            return max(prediction_time - self.observation_window * 24 * 3600, 0)
+        return 0
+
+    def transform(self, record: Dict[str, Any]) -> Dict[str, Any]:
+        person_id = record["person_id"]
+        prediction_times = self.person_index_date_map[person_id]
+        prediction_start_end_times = [
+            (
+                self._calculate_prediction_start_time(
+                    prediction_time_label_map["index_date"].timestamp()
+                ),
+                prediction_time_label_map["index_date"].timestamp(),
+                prediction_time_label_map["label"],
+            )
+            for prediction_time_label_map in prediction_times
+        ]
+        observation_window_indices = np.zeros(
+            (len(prediction_times), len(record["epoch_times"])), dtype=bool
+        )
+        for i, epoch_time in enumerate(record["epoch_times"]):
+            for sample_n, (
+                feature_extraction_time_start,
+                feature_extraction_end_end,
+                _,
+            ) in enumerate(prediction_start_end_times):
+                if (
+                    feature_extraction_time_start
+                    <= epoch_time
+                    <= feature_extraction_end_end
+                ):
+                    observation_window_indices[sample_n][i] = True
+
+        seq_length = len(record["epoch_times"])
+        time_series_columns = ["concept_ids", "input_ids"]
+        static_inputs = dict()
+        for k, v in record.items():
+            if k in ["concept_ids", "input_ids"]:
+                continue
+            if isinstance(v, (list, np.ndarray)) and len(v) == seq_length:
+                time_series_columns.append(k)
+            else:
+                static_inputs[k] = v
+
+        batched_samples = defaultdict(list)
+        for (_, index_date, label), observation_window_index in zip(
+            prediction_start_end_times, observation_window_indices
+        ):
+            for k, v in static_inputs.items():
+                batched_samples[k].append(v)
+            batched_samples["classifier_label"].append(label)
+            batched_samples["index_date"].append(index_date)
+            try:
+                start_age = int(record["concept_ids"][1].split(":")[1])
+            except Exception:
+                start_age = -1
+            batched_samples["age_at_index"].append(start_age)
+            for time_series_column in time_series_columns:
+                batched_samples[time_series_column].append(
+                    np.asarray(record[time_series_column])[observation_window_index]
+                )
+        return batched_samples
+
+    def batch_transform(self, record: Dict[str, Any]) -> Dict[str, Any]:
+        all_batched_record = defaultdict(list)
+        all_columns = record.keys()
+        for i in range(len(record["concept_ids"])):
+            one_record = {}
+            for column in all_columns:
+                one_record[column] = record[column][i]
+            new_batched_record = self.transform(one_record)
+            for k, v in new_batched_record.items():
+                all_batched_record[k].extend(v)
+        return all_batched_record

cehrgpt/data/sample_packing_sampler.py

@@ -1,5 +1,6 @@
 from typing import Iterator, List, Optional
 
+import numpy as np
 import torch
 import torch.distributed as dist
 from torch.utils.data import Sampler
@@ -33,6 +34,8 @@ class SamplePackingBatchSampler(Sampler[List[int]]):
         rank: Optional[int] = None,
         seed: int = 0,
         drop_last: bool = False,
+        negative_sampling_probability: Optional[float] = None,
+        labels: Optional[List[int]] = None,
     ):
         """
         Args:
@@ -73,6 +76,11 @@ class SamplePackingBatchSampler(Sampler[List[int]]):
                 f"Invalid rank {rank}, rank should be in the interval [0, {num_replicas - 1}]"
             )
 
+        if negative_sampling_probability is not None and labels is None:
+            raise ValueError(
+                f"When the negative sampling probability is provide, the labels must be provided as well"
+            )
+
         self.lengths = lengths
         self.max_tokens_per_batch = max_tokens_per_batch
         self.max_position_embeddings = max_position_embeddings
@@ -80,6 +88,8 @@ class SamplePackingBatchSampler(Sampler[List[int]]):
         self.rank = rank
         self.seed = seed
         self.drop_last = drop_last
+        self.negative_sampling_probability = negative_sampling_probability
+        self.labels = labels
         # Trainer https://github.com/huggingface/transformers/blame/main/src/transformers/trainer.py#L2470
         # http://github.com/huggingface/accelerate/blob/v0.31.0/src/accelerate/data_loader.py#L482
         # the huggingface trainer will call the accelerate.data_loader.DataLoaderShard.set_epoch,
@@ -100,6 +110,14 @@ class SamplePackingBatchSampler(Sampler[List[int]]):
         current_batch_tokens = 0
 
         for idx in indices:
+            # There is a chance to skip the negative samples to account for the class imbalance
+            # in the fine-tuning dataset
+            if self.negative_sampling_probability:
+                if (
+                    np.random.random() > self.negative_sampling_probability
+                    and self.labels[idx] == 0
+                ):
+                    continue
             # We take the minimum of the two because each sequence will be truncated to fit
             # the context window of the model
             sample_length = min(self.lengths[idx], self.max_position_embeddings)
@@ -131,10 +149,22 @@ class SamplePackingBatchSampler(Sampler[List[int]]):
         if len(self.lengths) == 0:
             return 0
 
-        # We need to truncate the lengths due to the context window limit imposed by the model
-        truncated_lengths = [
-            min(self.max_position_embeddings, length + 2) for length in self.lengths
-        ]
+        # There is a chance to skip the negative samples to account for the class imbalance
+        # in the fine-tuning dataset
+        if self.negative_sampling_probability:
+            truncated_lengths = []
+            for length, label in zip(self.lengths, self.labels):
+                if (
+                    np.random.random() > self.negative_sampling_probability
+                    and label == 0
+                ):
+                    continue
+                truncated_lengths.append(length)
+        else:
+            # We need to truncate the lengths due to the context window limit imposed by the model
+            truncated_lengths = [
+                min(self.max_position_embeddings, length + 2) for length in self.lengths
+            ]
 
         # Calculate average sequence length
         avg_seq_length = sum(truncated_lengths) // len(truncated_lengths)
@@ -145,7 +175,7 @@ class SamplePackingBatchSampler(Sampler[List[int]]):
         # Estimate total number of batches
         if self.drop_last:
             # If dropping last incomplete batch
-            return len(truncated_lengths) // seqs_per_batch * self.num_replicas
+            return len(truncated_lengths) // seqs_per_batch
         else:
             # If keeping last incomplete batch, ensure at least 1 batch
-            return max(1, len(truncated_lengths) // seqs_per_batch) * self.num_replicas
+            return max(1, len(truncated_lengths) // seqs_per_batch)

cehrgpt/generation/omop_converter_batch.py

@@ -60,6 +60,24 @@ OOV_CONCEPT_MAP = {
 }
 
 
+def extract_gender_concept_id(gender_token: str) -> int:
+    if gender_token.startswith("Gender/"):
+        return int(gender_token[len("Gender/") :])
+    elif gender_token.isnumeric():
+        return int(gender_token)
+    else:
+        return 0
+
+
+def extract_race_concept_id(race_token: str) -> int:
+    if race_token.startswith("Race/"):
+        return int(race_token[len("Race/") :])
+    elif race_token.isnumeric():
+        return int(race_token)
+    else:
+        return 0
+
+
 def create_folder_if_not_exists(output_folder, table_name):
     if not os.path.isdir(Path(output_folder) / table_name):
         os.mkdir(Path(output_folder) / table_name)
@@ -288,7 +306,13 @@ def gpt_to_omop_converter_batch(
         if int(birth_year) < 1900 or int(birth_year) > datetime.date.today().year:
             continue
 
-        p = Person(person_id, start_gender, birth_year, start_race)
+        p = Person(
+            person_id=person_id,
+            gender_concept_id=extract_gender_concept_id(start_gender),
+            year_of_birth=birth_year,
+            race_concept_id=extract_race_concept_id(start_race),
+        )
+
         append_to_dict(omop_export_dict, p, person_id)
         id_mappings_dict["person"][person_id] = person_id
         pt_seq_dict[person_id] = " ".join(concept_ids)
@@ -316,7 +340,12 @@ def gpt_to_omop_converter_batch(
                     id_mappings_dict["death"][person_id] = person_id
                 else:
                     try:
-                        visit_concept_id = int(clinical_events[event_idx + 1])
+                        if clinical_events[event_idx + 1].startswith("Visit/"):
+                            visit_concept_id = int(
+                                clinical_events[event_idx + 1][len("Visit/") :]
+                            )
+                        else:
+                            visit_concept_id = int(clinical_events[event_idx + 1])
                         inpatient_visit_indicator = visit_concept_id in [
                             9201,
                             262,
@@ -349,6 +378,7 @@ def gpt_to_omop_converter_batch(
                         visit_occurrence_id
                     ] = person_id
                     visit_occurrence_id += 1
+
             elif event in ATT_TIME_TOKENS:
                 if event[0] == "D":
                     att_date_delta = int(event[1:])

cehrgpt/gpt_utils.py

@@ -11,6 +11,7 @@ from cehrgpt.models.special_tokens import (
 )
 
 # Regular expression pattern to match inpatient attendance tokens
+MEDS_CODE_PATTERN = re.compile(r".*/.*")
 INPATIENT_ATT_PATTERN = re.compile(r"(?:VS-|i-)D(\d+)(?:-VE)?")
 DEMOGRAPHIC_PROMPT_SIZE = 4
 
@@ -194,8 +195,12 @@ def get_cehrgpt_output_folder(args, cehrgpt_tokenizer) -> str:
     return folder_name
 
 
-def is_clinical_event(token: str) -> bool:
-    return token.isnumeric()
+def is_clinical_event(token: str, meds: bool = False) -> bool:
+    if token.isnumeric():
+        return True
+    if meds:
+        return bool(MEDS_CODE_PATTERN.match(token))
+    return False
 
 
 def is_visit_start(token: str):
@@ -212,6 +217,18 @@ def is_visit_end(token: str) -> bool:
     return token in ["VE", "[VE]"]
 
 
+def is_inpatient_hour_token(token: str) -> bool:
+    return token.startswith("i-H")
+
+
+def extract_time_interval_in_hours(token: str) -> int:
+    try:
+        hour = int(token[3:])
+        return hour
+    except ValueError:
+        return 0
+
+
 def is_att_token(token: str):
     """
     Check if the token is an attention token.
@@ -251,6 +268,7 @@ def is_artificial_token(token: str) -> bool:
         return True
     if token == END_TOKEN:
         return True
+
     return False
 
 

cehrgpt/models/config.py

@@ -121,6 +121,7 @@ class CEHRGPTConfig(PretrainedConfig):
         bos_token_id=50256,
         eos_token_id=50256,
         lab_token_ids=None,
+        ve_token_id=None,
         scale_attn_by_inverse_layer_idx=False,
         reorder_and_upcast_attn=False,
         exclude_position_ids=False,
@@ -128,6 +129,10 @@ class CEHRGPTConfig(PretrainedConfig):
         value_vocab_size=None,
         include_ttv_prediction=False,
         use_sub_time_tokenization=True,
+        include_motor_time_to_event=True,
+        motor_tte_vocab_size=None,
+        motor_time_to_event_weight=1.0,
+        motor_num_time_pieces=16,
         token_to_time_token_mapping: Dict[int, List] = None,
         use_pretrained_embeddings=False,
         n_pretrained_embeddings_layers=2,
@@ -144,6 +149,7 @@ class CEHRGPTConfig(PretrainedConfig):
         entropy_penalty=False,
         entropy_penalty_alpha=0.01,
         sample_packing_max_positions=None,
+        class_weights=None,
         **kwargs,
     ):
         if token_to_time_token_mapping is None:
@@ -192,6 +198,22 @@ class CEHRGPTConfig(PretrainedConfig):
         self._token_to_time_token_mapping = token_to_time_token_mapping
         self.time_token_loss_weight = time_token_loss_weight
         self.time_to_visit_loss_weight = time_to_visit_loss_weight
+
+        # MOTOR TTE configuration
+        self.motor_tte_vocab_size = motor_tte_vocab_size
+        self.include_motor_time_to_event = (
+            include_motor_time_to_event
+            and self.motor_tte_vocab_size
+            and self.motor_tte_vocab_size > 0
+        )
+        if self.include_motor_time_to_event and not ve_token_id:
+            raise RuntimeError(
+                f"ve_token_id must be provided when include_motor_time_to_event is True"
+            )
+        self.ve_token_id = ve_token_id
+        self.motor_time_to_event_weight = motor_time_to_event_weight
+        self.motor_num_time_pieces = motor_num_time_pieces
+
         self.causal_sfm = causal_sfm
         self.demographics_size = demographics_size
         self.use_pretrained_embeddings = use_pretrained_embeddings
@@ -206,6 +228,9 @@ class CEHRGPTConfig(PretrainedConfig):
         self.entropy_penalty_alpha = entropy_penalty_alpha
         self.value_prediction_loss_weight = value_prediction_loss_weight
 
+        # Class weights for fine-tuning
+        self.class_weights = class_weights
+
         kwargs["tie_word_embeddings"] = not use_pretrained_embeddings
 
         super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)