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

cehrgpt/gpt_utils.py

@@ -1,7 +1,12 @@
 import random
 import re
-from datetime import date, timedelta
-from typing import List, Sequence, Tuple
+from datetime import date, datetime, timedelta, timezone
+from typing import List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+from cehrbert_data.const.artificial_tokens import DEATH_TOKEN
+from meds import death_code
+from transformers.utils import logging
 
 from cehrgpt.cehrgpt_args import SamplingStrategy
 from cehrgpt.models.special_tokens import (
@@ -14,6 +19,7 @@ from cehrgpt.models.special_tokens import (
 MEDS_CODE_PATTERN = re.compile(r".*/.*")
 INPATIENT_ATT_PATTERN = re.compile(r"(?:VS-|i-)D(\d+)(?:-VE)?")
 DEMOGRAPHIC_PROMPT_SIZE = 4
+logger = logging.get_logger("transformers")
 
 
 class RandomSampleCache:
@@ -62,6 +68,68 @@ class RandomSampleCache:
         return self._cache.pop()
 
 
+def construct_time_sequence(
+    concept_ids: List[str], epoch_times: Optional[List[Union[int, float]]] = None
+) -> List[float]:
+    if epoch_times is not None:
+        return epoch_times
+
+    if concept_ids[0].lower().startswith("year"):
+        year_str = concept_ids[0].split(":")[1]
+    else:
+        year_str = "1985"
+
+    datetime_cursor = datetime(
+        int(year_str), month=1, day=1, hour=0, minute=0, second=0
+    ).replace(tzinfo=timezone.utc)
+    epoch_times = []
+    for concept_id in concept_ids:
+        if is_att_token(concept_id):
+            att_days = extract_time_interval_in_days(concept_id)
+            datetime_cursor += timedelta(days=att_days)
+        epoch_times.append(datetime_cursor.timestamp())
+    return epoch_times
+
+
+def construct_age_sequence(
+    concept_ids: List[str], ages: Optional[List[int]] = None
+) -> List[int]:
+    if ages is not None:
+        return ages
+    elif concept_ids[1].lower().startswith("age"):
+        age_str = concept_ids[1].split(":")[1]
+        assert age_str.isnumeric(), f"age_str: {age_str}"
+        ages = []
+        time_delta = 0
+        for concept_id in concept_ids:
+            if is_att_token(concept_id):
+                time_delta += extract_time_interval_in_days(concept_id)
+            ages.append(int(age_str) + time_delta // 365)
+        return ages
+    else:
+        logger.warning(
+            "The second token is not a valid age token. The first 4 tokens are: %s. "
+            "Trying to fall back to ages, but it is not valid either %s. "
+            "Fall back to a zero vector [0, 0, 0, ...., 0]",
+            concept_ids[:4],
+            ages,
+        )
+        return np.zeros_like(concept_ids, dtype=int).tolist()
+
+
+def multiple_of_10(n: int) -> int:
+    return ((n // 10) + 1) * 10
+
+
+def encode_demographics(
+    age: int, gender: int, race: int, max_age=200, max_gender=10, max_race=10
+) -> int:
+    assert 0 <= age < max_age, f"age: {age}"
+    assert 0 <= gender < max_gender, f"gender: {gender}"
+    assert 0 <= race < max_race, f"race: {race}"
+    return age + max_age * gender + max_age * max_gender * race
+
+
 def collect_demographic_prompts_at_visits(patient_history: List[str]):
     demographic_prompts_at_visits = []
     start_year, start_age, start_gender, start_race = patient_history[
@@ -156,7 +224,7 @@ def random_slice_gpt_sequence(concept_ids, max_seq_len):
             )
         ):
             current_token = concept_ids[i]
-            if current_token == "VE":
+            if is_visit_end(current_token):
                 random_end_index = i
                 break
         return random_starting_index, random_end_index, demographic_tokens
@@ -198,6 +266,8 @@ def get_cehrgpt_output_folder(args, cehrgpt_tokenizer) -> str:
 def is_clinical_event(token: str, meds: bool = False) -> bool:
     if token.isnumeric():
         return True
+    if token in [DEATH_TOKEN, death_code]:
+        return True
     if meds:
         return bool(MEDS_CODE_PATTERN.match(token))
     return False

cehrgpt/models/activations.py

@@ -0,0 +1,27 @@
+# From https://github.com/bzhangGo/rmsnorm/blob/master/rmsnorm_torch.py
+# coding=utf-8
+
+from __future__ import absolute_import, division, print_function
+
+import torch
+import torch.nn as nn
+import transformers.pytorch_utils
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Mistral
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """MistralRMSNorm is equivalent to T5LayerNorm."""
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+transformers.pytorch_utils.ALL_LAYERNORM_LAYERS.extend([RMSNorm])

cehrgpt/models/config.py

@@ -106,6 +106,8 @@ class CEHRGPTConfig(PretrainedConfig):
         n_head=12,
         n_inner=None,
         activation_function="gelu_new",
+        decoder_mlp="GPT2MLP",
+        mlp_bias=False,
         resid_pdrop=0.1,
         embd_pdrop=0.1,
         attn_pdrop=0.1,
@@ -124,7 +126,7 @@ class CEHRGPTConfig(PretrainedConfig):
         ve_token_id=None,
         scale_attn_by_inverse_layer_idx=False,
         reorder_and_upcast_attn=False,
-        exclude_position_ids=False,
+        apply_rotary=False,
         include_values=False,
         value_vocab_size=None,
         include_ttv_prediction=False,
@@ -169,6 +171,8 @@ class CEHRGPTConfig(PretrainedConfig):
         self.n_head = n_head
         self.n_inner = n_inner
         self.activation_function = activation_function
+        self.decoder_mlp = decoder_mlp
+        self.mlp_bias = mlp_bias
         self.resid_pdrop = resid_pdrop
         self.embd_pdrop = embd_pdrop
         self.attn_pdrop = attn_pdrop
@@ -188,7 +192,7 @@ class CEHRGPTConfig(PretrainedConfig):
         self.eos_token_id = eos_token_id
         self.lab_token_ids = lab_token_ids
 
-        self.exclude_position_ids = exclude_position_ids
+        self.apply_rotary = apply_rotary
         self.include_values = include_values
         self.value_vocab_size = value_vocab_size