cehrgpt 0.0.1__py3-none-any.whl

cehrgpt/models/tokenization_hf_cehrgpt.py

@@ -0,0 +1,1077 @@
+import collections
+import copy
+import json
+import os
+import pickle
+from functools import partial
+from itertools import islice
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import scipy.stats as stats
+import transformers
+from cehrbert.models.hf_models.tokenization_utils import (
+    agg_helper,
+    agg_statistics,
+    load_json_file,
+)
+from cehrbert.runners.hf_runner_argument_dataclass import DataTrainingArguments
+from datasets import Dataset, DatasetDict
+from femr.stat_utils import OnlineStatistics, ReservoirSampler
+from scipy.interpolate import UnivariateSpline
+from tokenizers import AddedToken, Tokenizer
+from tokenizers.models import WordLevel
+from tokenizers.pre_tokenizers import WhitespaceSplit
+from tokenizers.trainers import WordLevelTrainer
+from tqdm import tqdm
+from transformers import PreTrainedTokenizer
+
+from cehrgpt.gpt_utils import (
+    convert_time_interval_to_time_tuple,
+    extract_time_interval_in_days,
+    is_att_token,
+    is_inpatient_att_token,
+)
+from cehrgpt.models.pretrained_embeddings import PretrainedEmbeddings
+from cehrgpt.models.special_tokens import (
+    END_TOKEN,
+    OUT_OF_VOCABULARY_TOKEN,
+    PAD_TOKEN,
+    START_TOKEN,
+)
+
+NUM_OF_BINS = 10
+DEGREE_OF_FREEDOM = 3
+SAMPLE_SIZE = 10_000
+NA = "N/A"
+UNKNOWN_BIN = "BIN:unknown"
+NONE_BIN = "BIN:NONE"
+TOKENIZER_FILE_NAME = "cehrgpt_tokenizer.json"
+VALUE_TOKENIZER_FILE_NAME = "cehrgpt_value_tokenizer.json"
+TIME_TOKENIZER_FILE_NAME = "cehrgpt_time_tokenizer.json"
+TOKEN_TO_SUB_TIME_TOKEN_MAPPING_FILE_NAME = "token_to_sub_time_token_mapping.json"
+LAB_STATS_FILE_NAME = "cehrgpt_lab_stats.pickle"
+LEGACY_LAB_STATS_FILE_NAME = "cehrgpt_lab_stats.json"
+CONCEPT_MAPPING_FILE_NAME = "concept_name_mapping.json"
+
+
+def truncated_sample(sample, standard_deviation):
+    lower_quantile = stats.norm.cdf(-standard_deviation)
+    upper_quantile = stats.norm.cdf(standard_deviation)
+    lower_bound = np.quantile(sample, lower_quantile)
+    upper_bound = np.quantile(sample, upper_quantile)
+    return [x for x in sample if lower_bound <= x <= upper_bound]
+
+
+def is_valid_valid_bin(token: str) -> bool:
+    return token.startswith("BIN:")
+
+
+def create_value_bin(bin_index: int) -> str:
+    return "BIN:" + str(bin_index)
+
+
+def create_sample_from_bins(bins, sample_size: int = 10_000) -> List[float]:
+    """
+    Generates a specified number of samples from a list of bins, each containing a fitted spline.
+
+    This function iterates over each bin, extracts the spline, and uses it to generate a set of samples
+    uniformly distributed along the x-axis defined by the spline's knots. It ensures that the total number
+    of samples generated matches the specified sample size by distributing the number of samples evenly
+    across the bins.
+
+    Parameters:
+        bins (List[Dict[str, UnivariateSpline]]): A list of dictionaries, each containing a 'spline' key
+            with a UnivariateSpline object as its value. These splines define the data distribution within
+            each bin from which samples are to be generated.
+        sample_size (int, optional): The total number of samples to generate from all bins combined.
+            Defaults to 10,000.
+
+    Returns:
+        List[float]: A list of sampled values, where each value is generated based on the spline functions
+            provided in the bins. The total number of samples in the list will be equal to `sample_size`.
+
+    Raises:
+        ValueError: If `sample_size` is less than the number of bins, as it would not be possible to generate
+            at least one sample per bin.
+
+    Example:
+        >>> x = np.linspace(0, 10, 100)
+        >>> y = np.sin(x)
+        >>> spline = UnivariateSpline(x, y, s=1)
+        >>> bins = [{'spline': spline} for _ in range(5)]
+        >>> samples = create_sample_from_bins(bins, 1000)
+        >>> len(samples)
+        1000
+
+    Note:
+        The function assumes that each bin's spline has a sufficient range of x-values (knots) to allow for
+        meaningful sampling. If the range of x-values is too narrow, the uniformity of the sample distribution
+        may be affected.
+    """
+    sample = []
+    num_of_bins = len(bins)
+    if num_of_bins > 0:
+        sample_per_bin = sample_size // num_of_bins
+        for value_bin in bins:
+            bin_spline = value_bin["spline"]
+            x = np.random.uniform(
+                bin_spline.get_knots()[0], bin_spline.get_knots()[-1], sample_per_bin
+            )
+            y = bin_spline(x)
+            sample.extend(y)
+    return sample
+
+
+def create_bins_with_spline(samples, num_bins, d_freedom=3) -> List[Dict[str, Any]]:
+    """
+    Divides a list of numeric samples into a specified number of bins and fits a spline to the data in each bin.
+
+    This function first sorts the list of samples, then partitions the sorted list into `num_bins` bins. For each bin,
+    a UnivariateSpline is fitted to the data within the bin, using the specified degrees of freedom. The function
+    handles edge cases by assigning infinity to the bounds of the first and last bins.
+
+    Parameters:
+        samples (List[float]): A list of sample data points, which are real numbers.
+        num_bins (int): The number of bins to divide the sample data into. It is assumed that there are enough samples to at least fill the bins to the minimum required for spline fitting.
+        d_freedom (int, optional): The degree of freedom for the spline. Default is 1, which fits a linear spline.
+
+    Returns:
+        List[Dict[str, Any]]: A list of dictionaries, each representing a bin. Each dictionary contains:
+            - 'bin_index' (int): The index of the bin.
+            - 'start_val' (float): The starting value of the bin, with the first bin starting at negative infinity.
+            - 'end_val' (float): The ending value of the bin, with the last bin ending at positive infinity.
+            - 'spline' (UnivariateSpline): The spline object fitted to the data within the bin.
+
+    Raises:
+        ValueError: If `num_bins` is less than 1 or if there are insufficient samples to create the specified number of bins with the required minimum data points per bin.
+
+    Example:
+        >>> samples = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+        >>> bins = create_bins_with_spline(samples, 2)
+        >>> for b in bins:
+        ...     print(b['bin_index'], b['start_val'], b['end_val'])
+        ...
+        0 -inf 5.5
+        1 5.5 inf
+
+    Note:
+        The function assumes that each bin will have at least `d_freedom + 1` samples to fit the spline. If the total number of samples is less than `num_bins * (d_freedom + 1)`, no bins will be created.
+    """
+    samples.sort()
+    bins = []
+    if len(samples) >= num_bins * (d_freedom + 1):
+        samples_per_bin = len(samples) // num_bins
+        for bin_index in range(0, num_bins):
+            if bin_index == 0:
+                start_val = float("-inf")
+            else:
+                start_val = samples[bin_index * samples_per_bin]
+
+            if bin_index == num_bins - 1:
+                end_val = float("inf")
+            else:
+                end_val = samples[(bin_index + 1) * samples_per_bin]
+            x = range(bin_index * samples_per_bin, (bin_index + 1) * samples_per_bin)
+            y = samples[bin_index * samples_per_bin : (bin_index + 1) * samples_per_bin]
+            spline = UnivariateSpline(x, y, k=d_freedom)
+            bins.append(
+                {
+                    "bin_index": bin_index,
+                    "start_val": start_val,
+                    "end_val": end_val,
+                    "spline": spline,
+                }
+            )
+    return bins
+
+
+def map_statistics(batch: Dict[str, Any], size=10_000) -> Dict[str, Any]:
+    if "units" in batch:
+        batch_value_units = batch["units"]
+    else:
+        batch_value_units = [[NA for _ in cons] for cons in batch["concept_ids"]]
+
+    if "number_as_values" not in batch:
+        batched_number_as_values = [
+            [value if isinstance(value, float) else None for value in concept_values]
+            for concept_values in batch["concept_values"]
+        ]
+    else:
+        batched_number_as_values = batch["number_as_values"]
+
+    if "concept_as_values" not in batch:
+        batched_concept_as_values = [
+            [value if isinstance(value, str) else None for value in concept_values]
+            for concept_values in batch["concept_values"]
+        ]
+    else:
+        batched_concept_as_values = batch["concept_as_values"]
+
+    numeric_stats_by_lab = collections.defaultdict(partial(ReservoirSampler, size=size))
+    categorical_stats_by_lab = collections.defaultdict(int)
+    for (
+        concept_ids,
+        number_as_values,
+        concept_as_values,
+        concept_value_indicators,
+        units,
+    ) in zip(
+        batch["concept_ids"],
+        batched_number_as_values,
+        batched_concept_as_values,
+        batch["concept_value_masks"],
+        batch_value_units,
+    ):
+        for (
+            concept_id,
+            number_as_value,
+            concept_as_value,
+            concept_value_indicator,
+            unit,
+        ) in zip(
+            concept_ids,
+            number_as_values,
+            concept_as_values,
+            concept_value_indicators,
+            units,
+        ):
+            if concept_value_indicator == 1:
+                if number_as_value:
+                    numeric_stats_by_lab[(concept_id, unit)].add(number_as_value, 1)
+                if concept_as_value:
+                    categorical_stats_by_lab[(concept_id, concept_as_value)] += 1
+
+    return {
+        "numeric_stats_by_lab": numeric_stats_by_lab,
+        "categorical_stats_by_lab": categorical_stats_by_lab,
+    }
+
+
+def create_numeric_concept_unit_mapping(
+    lab_stats: List[Dict[str, Any]]
+) -> Tuple[Dict[str, List[float]], Dict[str, List[str]]]:
+    numeric_concept_unit_mapping = collections.defaultdict(list)
+    for each_lab_stat in lab_stats:
+        numeric_concept_unit_mapping[each_lab_stat["concept_id"]].append(
+            (each_lab_stat["count"], each_lab_stat["unit"])
+        )
+
+    concept_prob_mapping = dict()
+    concept_unit_mapping = dict()
+    for concept_id in numeric_concept_unit_mapping.keys():
+        counts, units = zip(*numeric_concept_unit_mapping[concept_id])
+        total_count = sum(counts)
+        probs = [float(c) / total_count for c in counts]
+        concept_prob_mapping[concept_id] = probs
+        concept_unit_mapping[concept_id] = units
+    return concept_prob_mapping, concept_unit_mapping
+
+
+class NumericEventStatistics:
+    def __init__(self, lab_stats: List[Dict[str, Any]]):
+        self._lab_stats = lab_stats
+        self._lab_stats_mapping = {
+            (lab_stat["concept_id"], lab_stat["unit"]): {
+                "unit": lab_stat["unit"],
+                "mean": lab_stat["mean"],
+                "std": lab_stat["std"],
+                "value_outlier_std": lab_stat["value_outlier_std"],
+                "bins": lab_stat["bins"],
+            }
+            for lab_stat in lab_stats
+        }
+        self._concept_prob_mapping, self._concept_unit_mapping = (
+            create_numeric_concept_unit_mapping(lab_stats)
+        )
+
+    def get_numeric_concept_ids(self) -> List[str]:
+        return [_["concept_id"] for _ in self._lab_stats]
+
+    def get_random_unit(self, concept_id: str) -> str:
+        if concept_id in self._concept_prob_mapping:
+            unit_probs = self._concept_prob_mapping[concept_id]
+            return np.random.choice(
+                self._concept_unit_mapping[concept_id], p=unit_probs
+            )
+        return NA
+
+    def normalize(
+        self, concept_id: str, unit: str, concept_value: Union[float, str]
+    ) -> str:
+        if isinstance(concept_value, float):
+            if (concept_id, unit) in self._lab_stats_mapping:
+                concept_unit_stats = self._lab_stats_mapping[(concept_id, unit)]
+                bins = concept_unit_stats["bins"]
+                if bins:
+                    for each_bin in bins:
+                        if (
+                            each_bin["start_val"]
+                            <= concept_value
+                            <= each_bin["end_val"]
+                        ):
+                            return create_value_bin(each_bin["bin_index"])
+        return UNKNOWN_BIN
+
+    def denormalize(
+        self, concept_id: str, value_bin: str
+    ) -> Tuple[Optional[Union[float, str]], str]:
+        unit = self.get_random_unit(concept_id)
+        concept_value = value_bin
+        if (
+            is_valid_valid_bin(value_bin)
+            and (concept_id, unit) in self._lab_stats_mapping
+        ):
+            lab_stats = self._lab_stats_mapping[(concept_id, unit)]
+            bin_index = value_bin.split(":")[1]
+            if bin_index.isnumeric():
+                bin_index = int(bin_index)
+                # There are rare cases during sequence generation where bin_index could be out of range
+                # when there are no bins for (concept_id, unit) due to the small number of values in the source data
+                if len(lab_stats["bins"]) > bin_index:
+                    assert bin_index == lab_stats["bins"][bin_index]["bin_index"]
+                    bin_spline = lab_stats["bins"][bin_index]["spline"]
+                    x = np.random.uniform(
+                        bin_spline.get_knots()[0], bin_spline.get_knots()[-1]
+                    )
+                    concept_value = bin_spline(x).item()
+        return concept_value, unit
+
+
+class CehrGptTokenizer(PreTrainedTokenizer):
+
+    def __init__(
+        self,
+        tokenizer: Tokenizer,
+        value_tokenizer: Tokenizer,
+        att_tokenizer: Tokenizer,
+        token_to_sub_time_token_mapping: Dict[str, List[str]],
+        numeric_lab_stats: List[Dict[str, Any]],
+        categorical_lab_stats: Dict[Tuple[str, str], int],
+        concept_name_mapping: Dict[str, str],
+        pretrained_concept_embedding_model: PretrainedEmbeddings = None,
+    ):
+        self._tokenizer = tokenizer
+        self._value_tokenizer = value_tokenizer
+        self._att_tokenizer = att_tokenizer
+        self._token_to_sub_time_token_mapping = token_to_sub_time_token_mapping
+        self._numeric_lab_stats = numeric_lab_stats
+        self._numeric_event_statistics = NumericEventStatistics(numeric_lab_stats)
+        self._categorical_lab_stats = categorical_lab_stats
+        self._concept_name_mapping = concept_name_mapping
+        self._oov_token_id = self._tokenizer.token_to_id(OUT_OF_VOCABULARY_TOKEN)
+        self._padding_token_id = self._tokenizer.token_to_id(PAD_TOKEN)
+        self._start_token_id = self._tokenizer.token_to_id(START_TOKEN)
+        self._end_token_id = self._tokenizer.token_to_id(END_TOKEN)
+        self._numeric_concept_ids = (
+            self._numeric_event_statistics.get_numeric_concept_ids()
+        )
+        self._categorical_concept_ids = list(
+            {t[0] for t in self._categorical_lab_stats.keys()}
+        )
+        self._padding_value_token_id = self._value_tokenizer.token_to_id(PAD_TOKEN)
+        self._pretrained_concept_embedding_model = (
+            pretrained_concept_embedding_model
+            if pretrained_concept_embedding_model
+            else PretrainedEmbeddings(None)
+        )
+        self._pretrained_concept_ids = [
+            _
+            for _ in self.get_vocab().keys()
+            if self._pretrained_concept_embedding_model.is_concept_available(_)
+        ]
+
+        super().__init__()
+
+    @property
+    def pretrained_concept_ids(self):
+        return self._pretrained_concept_ids
+
+    @property
+    def pretrained_token_ids(self):
+        return self.encode(self._pretrained_concept_ids)
+
+    @property
+    def pretrained_embeddings(self):
+        return np.asarray(
+            [
+                self._pretrained_concept_embedding_model.get_concept_embeddings(_)
+                for _ in self._pretrained_concept_ids
+            ]
+        )
+
+    @property
+    def vocab_size(self) -> int:
+        return self._tokenizer.get_vocab_size()
+
+    @property
+    def value_vocab_size(self) -> int:
+        return self._value_tokenizer.get_vocab_size()
+
+    @property
+    def time_token_vocab_size(self) -> int:
+        return self._att_tokenizer.get_vocab_size()
+
+    @property
+    def pad_value_token_id(self):
+        return self._padding_value_token_id
+
+    @property
+    def start_token_id(self):
+        return self._start_token_id
+
+    @property
+    def end_token_id(self):
+        return self._end_token_id
+
+    @property
+    def end_token(self):
+        return END_TOKEN
+
+    @property
+    def eos_token(self):
+        return END_TOKEN
+
+    @property
+    def eos_token_id(self):
+        return self._end_token_id
+
+    @property
+    def pad_token_id(self):
+        return self._padding_token_id
+
+    @property
+    def pad_token(self):
+        return PAD_TOKEN
+
+    @property
+    def numeric_concept_ids(self):
+        return self._numeric_concept_ids
+
+    @property
+    def categorical_concept_ids(self):
+        return self._categorical_concept_ids
+
+    @property
+    def lab_token_ids(self):
+        reserved_tokens = [START_TOKEN, PAD_TOKEN, END_TOKEN, OUT_OF_VOCABULARY_TOKEN]
+        return self.encode(
+            [
+                concept_id
+                for concept_id in self._numeric_concept_ids
+                + self._categorical_concept_ids
+                if concept_id not in reserved_tokens
+            ]
+        )
+
+    @property
+    def token_to_time_token_mapping(self) -> Dict[int, List[int]]:
+        default_mapping = {-1: [0, 0, 0]}
+        default_mapping.update(
+            {
+                self._tokenizer.token_to_id(time_token): list(
+                    map(self._att_tokenizer.token_to_id, sub_time_tokens)
+                )
+                for time_token, sub_time_tokens in self._token_to_sub_time_token_mapping.items()
+            }
+        )
+        return default_mapping
+
+    @property
+    def pretrained_concept_embedding_model(self):
+        return self._pretrained_concept_embedding_model
+
+    def get_vocab(self) -> Dict[str, int]:
+        return self._tokenizer.get_vocab()
+
+    def get_value_vocab(self) -> Dict[str, int]:
+        return self._value_tokenizer.get_vocab()
+
+    def encode(self, concept_ids, **kwargs) -> Sequence[int]:
+        encoded = self._tokenizer.encode(concept_ids, is_pretokenized=True)
+        return encoded.ids
+
+    def decode(
+        self, concept_token_ids: List[int], skip_special_tokens: bool = True, **kwargs
+    ) -> List[str]:
+        return self._tokenizer.decode(
+            concept_token_ids, skip_special_tokens=skip_special_tokens
+        ).split(" ")
+
+    def encode_value(self, concept_values: Sequence[str]) -> Sequence[int]:
+        encoded = self._value_tokenizer.encode(concept_values, is_pretokenized=True)
+        return encoded.ids
+
+    def decode_value(
+        self, concept_value_token_ids: List[int], skip_special_tokens: bool = True
+    ) -> List[str]:
+        return self._value_tokenizer.decode(
+            concept_value_token_ids, skip_special_tokens=skip_special_tokens
+        ).split(" ")
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        token_id = self._tokenizer.token_to_id(token)
+        return token_id if token_id else self._oov_token_id
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        token = self._tokenizer.id_to_token(index)
+        return token if token else OUT_OF_VOCABULARY_TOKEN
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) in a single string."""
+        out_string = " ".join([self._concept_name_mapping[t] for t in tokens])
+        return out_string
+
+    def save_pretrained(
+        self,
+        save_directory: Union[str, os.PathLike],
+        push_to_hub: bool = False,
+        **kwargs,
+    ):
+        """
+        Save the Cehrbert tokenizer.
+
+        This method make sure the batch processor can then be re-loaded using the
+        .from_pretrained class method.
+
+        Args:
+            save_directory (`str` or `os.PathLike`): The path to a directory where the tokenizer will be saved.
+            push_to_hub (`bool`, *optional*, defaults to `False`):
+                Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
+                repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
+                namespace).
+            kwargs (`Dict[str, Any]`, *optional*):
+                Additional key word arguments passed along to the [`PushToHubMixin.push_to_hub`] method.
+        """
+        assert not os.path.isfile(
+            save_directory
+        ), f"Provided path ({save_directory}) should be a directory, not a file"
+
+        os.makedirs(save_directory, exist_ok=True)
+
+        if push_to_hub:
+            commit_message = kwargs.pop("commit_message", None)
+            repo_id = kwargs.pop("repo_id", str(save_directory).split(os.path.sep)[-1])
+            repo_id = self._create_repo(repo_id, **kwargs)
+            files_timestamps = self._get_files_timestamps(save_directory)
+
+        self._tokenizer.save(os.path.join(save_directory, TOKENIZER_FILE_NAME))
+
+        self._value_tokenizer.save(
+            os.path.join(save_directory, VALUE_TOKENIZER_FILE_NAME)
+        )
+
+        self._att_tokenizer.save(os.path.join(save_directory, TIME_TOKENIZER_FILE_NAME))
+
+        with open(
+            os.path.join(save_directory, TOKEN_TO_SUB_TIME_TOKEN_MAPPING_FILE_NAME), "w"
+        ) as f:
+            json.dump(self._token_to_sub_time_token_mapping, f)
+
+        with open(os.path.join(save_directory, LAB_STATS_FILE_NAME), "wb") as f:
+            lab_stats = {
+                "numeric_lab_stats": self._numeric_lab_stats,
+                "categorical_lab_stats": self._categorical_lab_stats,
+            }
+            pickle.dump(lab_stats, f)
+
+        with open(os.path.join(save_directory, CONCEPT_MAPPING_FILE_NAME), "w") as f:
+            json.dump(self._concept_name_mapping, f)
+
+        self._pretrained_concept_embedding_model.save(save_directory)
+
+        if push_to_hub:
+            self._upload_modified_files(
+                save_directory,
+                repo_id,
+                files_timestamps,
+                commit_message=commit_message,
+                token=kwargs.get("token"),
+            )
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Union[str, os.PathLike],
+        **kwargs,
+    ):
+        """
+        Load the CehrBert tokenizer.
+
+        Parameters:
+            pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
+                Can be either:
+                    - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
+                      Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced under a
+                      user or organization name, like `dbmdz/bert-base-german-cased`.
+                    - A path to a *directory* containing tokenization data saved using
+                      [`save_pretrained`], e.g., `./my_data_directory/`.
+            kwargs: Arguments for loading to pass to transformers.utils.hub.cached_file
+
+        Returns:
+            A CehrBert Tokenizer
+        """
+
+        is_legacy_tokenizer = CehrGptTokenizer.is_legacy_tokenizer(
+            pretrained_model_name_or_path, **kwargs
+        )
+
+        # Load the concept tokenizer
+        tokenizer_file = transformers.utils.hub.cached_file(
+            pretrained_model_name_or_path, TOKENIZER_FILE_NAME, **kwargs
+        )
+        if not tokenizer_file:
+            return None
+        tokenizer = Tokenizer.from_file(tokenizer_file)
+
+        # Load the concept_value_tokenizer
+        if is_legacy_tokenizer:
+            value_tokenizer = Tokenizer(
+                WordLevel(unk_token=OUT_OF_VOCABULARY_TOKEN, vocab=dict())
+            )
+        else:
+            value_tokenizer_file = transformers.utils.hub.cached_file(
+                pretrained_model_name_or_path, VALUE_TOKENIZER_FILE_NAME, **kwargs
+            )
+            if not value_tokenizer_file:
+                return None
+            value_tokenizer = Tokenizer.from_file(value_tokenizer_file)
+
+        # Load the ttt tokenizer
+        att_tokenizer_file = transformers.utils.hub.cached_file(
+            pretrained_model_name_or_path, TIME_TOKENIZER_FILE_NAME, **kwargs
+        )
+        if not att_tokenizer_file:
+            return None
+        att_tokenizer = Tokenizer.from_file(att_tokenizer_file)
+
+        # Load the sub time token json file
+        token_to_sub_time_token_mapping_file = transformers.utils.hub.cached_file(
+            pretrained_model_name_or_path,
+            TOKEN_TO_SUB_TIME_TOKEN_MAPPING_FILE_NAME,
+            **kwargs,
+        )
+        if not token_to_sub_time_token_mapping_file:
+            return None
+        token_to_sub_time_token_mapping = load_json_file(
+            token_to_sub_time_token_mapping_file
+        )
+
+        # Load the lab stats pickle file
+        if is_legacy_tokenizer:
+            legacy_lab_stats_file = transformers.utils.hub.cached_file(
+                pretrained_model_name_or_path, LEGACY_LAB_STATS_FILE_NAME, **kwargs
+            )
+            if not legacy_lab_stats_file:
+                return None
+            # Support the old version of the numeric lab stats file
+            lab_stats = {
+                "numeric_lab_stats": load_json_file(legacy_lab_stats_file),
+                "categorical_lab_stats": dict(),
+            }
+        else:
+            lab_stats_file = transformers.utils.hub.cached_file(
+                pretrained_model_name_or_path, LAB_STATS_FILE_NAME, **kwargs
+            )
+            if not lab_stats_file:
+                return None
+
+            with open(lab_stats_file, "rb") as file:
+                lab_stats = pickle.load(file)
+
+        # Load the concept_name json file
+        concept_name_mapping_file = transformers.utils.hub.cached_file(
+            pretrained_model_name_or_path, CONCEPT_MAPPING_FILE_NAME, **kwargs
+        )
+        if not concept_name_mapping_file:
+            return None
+        concept_name_mapping = load_json_file(concept_name_mapping_file)
+
+        pretrained_embedding_model = PretrainedEmbeddings(pretrained_model_name_or_path)
+
+        return CehrGptTokenizer(
+            tokenizer,
+            value_tokenizer,
+            att_tokenizer,
+            token_to_sub_time_token_mapping,
+            lab_stats["numeric_lab_stats"],
+            lab_stats["categorical_lab_stats"],
+            concept_name_mapping,
+            pretrained_embedding_model,
+        )
+
+    @classmethod
+    def is_legacy_tokenizer(
+        cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
+    ):
+        try:
+            legacy_lab_stats_file = transformers.utils.hub.cached_file(
+                pretrained_model_name_or_path, LEGACY_LAB_STATS_FILE_NAME, **kwargs
+            )
+            return legacy_lab_stats_file is not None
+        except Exception:
+            return False
+
+    @classmethod
+    def expand_trained_tokenizer(
+        cls,
+        cehrgpt_tokenizer,
+        dataset: Union[Dataset, DatasetDict],
+        concept_name_mapping: Dict[str, str],
+        data_args: DataTrainingArguments,
+        pretrained_concept_embedding_model: PretrainedEmbeddings = None,
+    ):
+        if not isinstance(cehrgpt_tokenizer, CehrGptTokenizer):
+            raise ValueError(
+                "The existing cehrgpt must be an instance of CehrGptTokenizer"
+            )
+
+        cehrgpt_tokenizer_copy = copy.deepcopy(cehrgpt_tokenizer)
+
+        new_tokenizer = CehrGptTokenizer.train_tokenizer(
+            dataset=dataset,
+            concept_name_mapping=concept_name_mapping,
+            data_args=data_args,
+        )
+
+        new_tokens = set(new_tokenizer.get_vocab().keys()) - set(
+            cehrgpt_tokenizer_copy.get_vocab().keys()
+        )
+        new_value_tokens = set(new_tokenizer.get_value_vocab().keys()) - set(
+            cehrgpt_tokenizer_copy.get_value_vocab().keys()
+        )
+        new_att_tokens = set(new_tokenizer._att_tokenizer.get_vocab().keys()) - set(
+            cehrgpt_tokenizer_copy._att_tokenizer.get_vocab().keys()
+        )
+        new_token_to_sub_time_token_mapping = (
+            new_tokenizer._token_to_sub_time_token_mapping
+        )
+        new_numeric_lab_stats = new_tokenizer._numeric_lab_stats
+        new_categorical_lab_stats = new_tokenizer._categorical_lab_stats
+        new_concept_name_mapping = new_tokenizer._concept_name_mapping
+
+        # Add new tokens to the existing tokenizer
+        cehrgpt_tokenizer_copy._tokenizer.add_tokens(
+            [
+                AddedToken(token, single_word=True, normalized=False)
+                for token in new_tokens
+            ]
+        )
+        # Add new tokens to the existing value tokenizer
+        cehrgpt_tokenizer_copy._value_tokenizer.add_tokens(
+            [
+                AddedToken(token, single_word=True, normalized=False)
+                for token in new_value_tokens
+            ]
+        )
+        # Add new time tokens to the existing att tokenizer
+        cehrgpt_tokenizer_copy._att_tokenizer.add_tokens(
+            [
+                AddedToken(token, single_word=True, normalized=False)
+                for token in new_att_tokens
+            ]
+        )
+        # Merge the time_token -> List[sub_time_tokens] mapping
+        for time_token, sub_time_tokens in new_token_to_sub_time_token_mapping.items():
+            if (
+                time_token
+                not in cehrgpt_tokenizer_copy._token_to_sub_time_token_mapping
+            ):
+                cehrgpt_tokenizer_copy._token_to_sub_time_token_mapping[time_token] = (
+                    sub_time_tokens
+                )
+
+        # Merge numeric lab_stats
+        cehrgpt_tokenizer_copy._numeric_lab_stats = cls.merge_numeric_lab_stats(
+            cehrgpt_tokenizer_copy._numeric_lab_stats,
+            new_numeric_lab_stats,
+        )
+        # Merge categorical lab_stats
+        cehrgpt_tokenizer_copy._categorical_lab_stats = cls.merge_categorical_lab_stats(
+            cehrgpt_tokenizer_copy._categorical_lab_stats,
+            new_categorical_lab_stats,
+        )
+
+        # Merge concept_name_mapping
+        for token, concept_name in new_concept_name_mapping.items():
+            if token not in cehrgpt_tokenizer_copy._concept_name_mapping:
+                cehrgpt_tokenizer_copy._concept_name_mapping[token] = concept_name
+
+        return CehrGptTokenizer(
+            tokenizer=cehrgpt_tokenizer_copy._tokenizer,
+            value_tokenizer=cehrgpt_tokenizer_copy._value_tokenizer,
+            att_tokenizer=cehrgpt_tokenizer_copy._att_tokenizer,
+            token_to_sub_time_token_mapping=cehrgpt_tokenizer_copy._token_to_sub_time_token_mapping,
+            numeric_lab_stats=cehrgpt_tokenizer_copy._numeric_lab_stats,
+            categorical_lab_stats=cehrgpt_tokenizer_copy._categorical_lab_stats,
+            concept_name_mapping=cehrgpt_tokenizer_copy._concept_name_mapping,
+            pretrained_concept_embedding_model=pretrained_concept_embedding_model,
+        )
+
+    @classmethod
+    def merge_numeric_lab_stats(
+        cls,
+        lab_stats_existing: List[Dict[str, Any]],
+        lab_stats_new: List[Dict[str, Any]],
+    ) -> List[Dict[str, Any]]:
+
+        lab_stats_existing_mapping = {
+            (lab_stat["concept_id"], lab_stat["unit"]): lab_stat
+            for lab_stat in lab_stats_existing
+        }
+        for lab_stat in lab_stats_new:
+            concept_unit_pair = (lab_stat["concept_id"], lab_stat["unit"])
+            if concept_unit_pair in lab_stats_existing_mapping:
+                existing = OnlineStatistics()
+                existing.count = lab_stats_existing_mapping[concept_unit_pair]["count"]
+                existing.current_mean = lab_stats_existing_mapping[concept_unit_pair][
+                    "mean"
+                ]
+                existing.variance = (
+                    lab_stats_existing_mapping[concept_unit_pair]["std"] ** 2
+                    * existing.count
+                )
+                new = OnlineStatistics()
+                new.count = lab_stat["count"]
+                new.current_mean = lab_stat["mean"]
+                new.variance = lab_stat["std"] ** 2 * new.count
+                existing.combine(new)
+                lab_stats_existing_mapping[concept_unit_pair]["mean"] = existing.mean()
+                lab_stats_existing_mapping[concept_unit_pair][
+                    "std"
+                ] = existing.standard_deviation()
+                lab_stats_existing_mapping[concept_unit_pair]["count"] = existing.count
+                # recreate the bins
+                sample = create_sample_from_bins(
+                    lab_stats_existing_mapping[concept_unit_pair]["bins"]
+                )
+                sample.extend(create_sample_from_bins(lab_stat["bins"]))
+                lab_stats_existing_mapping[concept_unit_pair]["bins"] = (
+                    create_bins_with_spline(sample, NUM_OF_BINS, DEGREE_OF_FREEDOM)
+                )
+
+            else:
+                if lab_stat["count"] > 0:
+                    lab_stats_existing_mapping[concept_unit_pair] = lab_stat
+
+        return list(lab_stats_existing_mapping.values())
+
+    @classmethod
+    def merge_categorical_lab_stats(
+        cls,
+        categorical_lab_stats_existing: Dict[Tuple[str, str], int],
+        categorical_lab_stats_new: Dict[Tuple[str, str], int],
+    ) -> Dict[Tuple[str, str], int]:
+        for (concept_id, concept_as_value), count in categorical_lab_stats_new.items():
+            if (concept_id, concept_as_value) not in categorical_lab_stats_new:
+                categorical_lab_stats_existing[(concept_id, concept_as_value)] = 0
+            categorical_lab_stats_existing[(concept_id, concept_as_value)] += count
+        return categorical_lab_stats_existing
+
+    @classmethod
+    def train_tokenizer(
+        cls,
+        dataset: Union[Dataset, DatasetDict],
+        concept_name_mapping: Dict[str, str],
+        data_args: DataTrainingArguments,
+        pretrained_concept_embedding_model: PretrainedEmbeddings = None,
+    ):
+        """
+        Train a huggingface word level tokenizer.
+
+        To use their tokenizer, we need to concatenate all the concepts
+        together and treat it as a sequence.
+        """
+
+        if isinstance(dataset, DatasetDict):
+            dataset = dataset["train"]
+
+        concept_tokenizer = cls.train_concept_tokenizer(
+            dataset,
+            feature_name="concept_ids",
+            special_tokens=[PAD_TOKEN, OUT_OF_VOCABULARY_TOKEN, START_TOKEN, END_TOKEN],
+            unk_token=OUT_OF_VOCABULARY_TOKEN,
+            data_args=data_args,
+        )
+        concept_value_column = "concept_as_values"
+        for row in dataset:
+            if concept_value_column not in row:
+                concept_value_column = "concept_values"
+            break
+        value_tokenizer = cls.train_concept_tokenizer(
+            dataset,
+            feature_name=concept_value_column,
+            special_tokens=[OUT_OF_VOCABULARY_TOKEN, PAD_TOKEN],
+            unk_token=OUT_OF_VOCABULARY_TOKEN,
+            data_args=data_args,
+        )
+        value_tokenizer.add_tokens(
+            [
+                AddedToken(_, single_word=True, normalized=False)
+                for _ in [create_value_bin(_) for _ in range(NUM_OF_BINS)]
+                + [UNKNOWN_BIN, NONE_BIN]
+            ]
+        )
+
+        map_statistics_partial = partial(map_statistics, size=SAMPLE_SIZE)
+
+        if data_args.streaming:
+            parts = dataset.map(
+                partial(agg_helper, map_func=map_statistics_partial),
+                batched=True,
+                batch_size=data_args.preprocessing_batch_size,
+                new_fingerprint="invalid",
+                remove_columns=dataset.column_names,
+            )
+        else:
+            parts = dataset.map(
+                partial(agg_helper, map_func=map_statistics_partial),
+                batched=True,
+                batch_size=data_args.preprocessing_batch_size,
+                remove_columns=dataset.column_names,
+                num_proc=data_args.preprocessing_num_workers,
+                keep_in_memory=True,
+                new_fingerprint="invalid",
+            )
+        current = None
+        for stat in tqdm(parts, desc="Aggregating the lab statistics"):
+            fixed_stat = pickle.loads(stat["data"])
+            if current is None:
+                current = fixed_stat
+            else:
+                current = agg_statistics(current, fixed_stat)
+
+        numeric_lab_stats = []
+        for (concept_id, unit), online_stats in current["numeric_stats_by_lab"].items():
+            if len(online_stats.samples) == 0:
+                continue
+            samples = truncated_sample(
+                online_stats.samples, data_args.value_outlier_std
+            )
+            bins = create_bins_with_spline(samples, NUM_OF_BINS, DEGREE_OF_FREEDOM)
+            if len(bins) > 0:
+                numeric_lab_stats.append(
+                    {
+                        "concept_id": concept_id,
+                        "unit": unit,
+                        "mean": np.mean(samples),
+                        "std": np.std(samples),
+                        "count": len(online_stats.samples),
+                        "value_outlier_std": data_args.value_outlier_std,
+                        "bins": bins,
+                    }
+                )
+
+        categorical_lab_stats = collections.defaultdict(int)
+        for (concept_id, value_as_concept), count in current[
+            "categorical_stats_by_lab"
+        ].items():
+            categorical_lab_stats[(concept_id, value_as_concept)] += count
+
+        # We will train a tokenizer specifically for time intervals
+        sub_time_token_data = []
+        token_to_sub_time_token_mapping = collections.defaultdict(list)
+        for token, token_id in concept_tokenizer.get_vocab().items():
+            if is_att_token(token):
+                time_interval = extract_time_interval_in_days(token)
+                time_tuple = convert_time_interval_to_time_tuple(
+                    time_interval, is_inpatient_att_token(token)
+                )
+                token_to_sub_time_token_mapping[token] = list(time_tuple)
+                sub_time_token_data.append(" ".join(time_tuple))
+
+        att_tokenizer = Tokenizer(
+            WordLevel(unk_token=OUT_OF_VOCABULARY_TOKEN, vocab=dict())
+        )
+        att_tokenizer.pre_tokenizer = WhitespaceSplit()
+        att_trainer = WordLevelTrainer(
+            special_tokens=[OUT_OF_VOCABULARY_TOKEN],
+            vocab_size=data_args.vocab_size,
+            min_frequency=0,
+            show_progress=True,
+        )
+        att_tokenizer.train_from_iterator(sub_time_token_data, trainer=att_trainer)
+
+        return CehrGptTokenizer(
+            concept_tokenizer,
+            value_tokenizer,
+            att_tokenizer,
+            token_to_sub_time_token_mapping,
+            numeric_lab_stats,
+            categorical_lab_stats,
+            concept_name_mapping,
+            pretrained_concept_embedding_model,
+        )
+
+    @classmethod
+    def train_concept_tokenizer(
+        cls,
+        dataset,
+        feature_name,
+        special_tokens: List[str],
+        unk_token,
+        data_args,
+    ):
+        # Use the Fast Tokenizer from the Huggingface tokenizers Rust implementation.
+        # https://github.com/huggingface/tokenizers
+        concept_tokenizer = Tokenizer(WordLevel(unk_token=unk_token, vocab=dict()))
+        concept_tokenizer.pre_tokenizer = WhitespaceSplit()
+        concept_trainer = WordLevelTrainer(
+            special_tokens=special_tokens,
+            vocab_size=data_args.vocab_size,
+            min_frequency=data_args.min_frequency,
+            show_progress=True,
+        )
+        batch_concat_concepts_partial_func = partial(
+            cls.batch_concat_concepts, feature_name=feature_name
+        )
+        if data_args.streaming:
+            concatenated_features = dataset.map(
+                batch_concat_concepts_partial_func,
+                batched=True,
+                batch_size=data_args.preprocessing_batch_size,
+            )
+
+            def batched_generator():
+                iterator = iter(concatenated_features)
+                while True:
+                    batch = list(islice(iterator, data_args.preprocessing_batch_size))
+                    if not batch:
+                        break
+                    yield [example[feature_name] for example in batch]
+
+            # We pass a generator of list of texts (concatenated concept_ids) to train_from_iterator
+            # for efficient training
+            generator = batched_generator()
+        else:
+            concatenated_features = dataset.map(
+                batch_concat_concepts_partial_func,
+                num_proc=data_args.preprocessing_num_workers,
+                batched=True,
+                batch_size=data_args.preprocessing_batch_size,
+                remove_columns=dataset.column_names,
+            )
+            generator = concatenated_features[feature_name]
+        concept_tokenizer.train_from_iterator(generator, trainer=concept_trainer)
+        return concept_tokenizer
+
+    def normalize(self, concept_id: str, unit: str, concept_value: float) -> str:
+        return self._numeric_event_statistics.normalize(concept_id, unit, concept_value)
+
+    def denormalize(self, concept_id: str, value_bin: str) -> Tuple[float, str]:
+        return self._numeric_event_statistics.denormalize(concept_id, value_bin)
+
+    @classmethod
+    def batch_concat_concepts(
+        cls, records: Dict[str, List], feature_name
+    ) -> Dict[str, List]:
+        return {
+            feature_name: [
+                " ".join(
+                    [token for token in tokens if token and isinstance(token, str)]
+                )
+                for tokens in records[feature_name]
+            ]
+        }