cehrgpt 0.0.1__py3-none-any.whl

cehrgpt/models/hf_modeling_outputs.py

@@ -0,0 +1,158 @@
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+from transformers.modeling_outputs import ModelOutput
+
+
+@dataclass
+class CehrGptOutputWithPast(ModelOutput):
+    """
+    Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding).
+
+    Args:
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+
+            If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1,
+            hidden_size)` is output.
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+            `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if
+            `config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads,
+            encoder_sequence_length, embed_size_per_head)`.
+
+            Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if
+            `config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values`
+            input) to speed up sequential decoding.
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    last_hidden_state: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class CehrGptCausalLMOutput(ModelOutput):
+    """
+    Base class for causal language model (or autoregressive) outputs.
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Language modeling loss (for next-token prediction).
+        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+            Tuple of `torch.FloatTensor` tuples of length `config.n_layers`, with each tuple containing the cached key,
+            value states of the self-attention and the cross-attention layers if model is used in encoder-decoder
+            setting. Only relevant if `config.is_decoder = True`.
+
+            Contains pre-computed hidden-states (key and values in the attention blocks) that can be used (see
+            `past_key_values` input) to speed up sequential decoding.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    value_indicators: torch.BoolTensor = None
+    next_value_logits: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+    token_loss: Optional[torch.FloatTensor] = None
+    time_token_loss: Optional[torch.FloatTensor] = None
+    time_to_visit_loss: Optional[torch.FloatTensor] = None
+    token_value_loss: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class CehrGptGenerateDecoderOnlyOutput(ModelOutput):
+    """
+    Outputs of decoder-only generation models, when using non-beam methods.
+
+    Args:
+        sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter
+            if all batches finished early due to the `eos_token_id`.
+        scores (`tuple(torch.FloatTensor)` *optional*, returned when `output_scores=True` is passed or when `config.output_scores=True`):
+            Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
+            at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for
+            each generated token), with each tensor of shape `(batch_size, config.vocab_size)`.
+        logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True` is passed or when `config.output_logits=True`):
+            Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
+            at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for
+            each generated token), with each tensor of shape `(batch_size, config.vocab_size)`.
+        attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True` is passed or `config.output_attentions=True`):
+            Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+            `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.
+        hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+            `torch.FloatTensor` of shape `(batch_size, generated_length, hidden_size)`.
+        past_key_values (`tuple(tuple(torch.FloatTensor)))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+            NOTE: some models have a different `past_key_values` format, confirm with the model's documentation.
+            Usually a Tuple (one element for each layer of the decoder) of tuples (two elements, key tensor and value
+            tensor). The first Tuple is of length `config.n_layers`, with each tuple having 2 tensors of shape
+            `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if
+            `config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads,
+            encoder_sequence_length, embed_size_per_head)`.
+    """
+
+    sequences: torch.LongTensor = None
+    sequence_val_masks: Optional[torch.BoolTensor] = None
+    sequence_vals: Optional[torch.LongTensor] = None
+    scores: Optional[Tuple[torch.FloatTensor]] = None
+    logits: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    past_key_values: Optional[Tuple[Tuple[Tuple[torch.FloatTensor]]]] = None
+
+
+@dataclass
+class CehrGptSequenceClassifierOutput(ModelOutput):
+    """
+    Base class for outputs of sentence classification models.
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Classification (or regression if config.num_labels==1) loss.
+        logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
+            Classification (or regression if config.num_labels==1) scores (before SoftMax).
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+        attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+            Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+            sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None

cehrgpt/models/pretrained_embeddings.py

@@ -0,0 +1,82 @@
+import pickle
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+
+PRETRAINED_EMBEDDING_VECTOR_FILE_NAME = "pretrained_embedding_vectors.npy"
+PRETRAINED_EMBEDDING_CONCEPT_FILE_NAME = "pretrained_embedding_concepts.pkl"
+
+
+class PretrainedEmbeddings:
+    """A class to handle pretrained embedding vectors and their associated concepts."""
+
+    def __init__(self, model_folder: Optional[str]):
+        if model_folder:
+            model_path = Path(model_folder)
+            self.vector_file = model_path / PRETRAINED_EMBEDDING_VECTOR_FILE_NAME
+            self.concept_file = model_path / PRETRAINED_EMBEDDING_CONCEPT_FILE_NAME
+            self.exists = self.vector_file.exists() and self.concept_file.exists()
+        else:
+            self.exists = False
+        self._initialize_embeddings() if self.exists else self._initialize_empty()
+
+    def _initialize_embeddings(self):
+        """Load embeddings and associated concepts from files."""
+        self.pretrained_embeddings = np.load(self.vector_file)
+        with open(self.concept_file, "rb") as f:
+            self.pretrained_concepts = pickle.load(f)
+
+        self.concept_ids = [
+            concept["concept_id"] for concept in self.pretrained_concepts
+        ]
+        self.reverse_concept_id_map = {
+            concept_id: i for i, concept_id in enumerate(self.concept_ids)
+        }
+        self.concept_names = [
+            concept["concept_name"] for concept in self.pretrained_concepts
+        ]
+        self.embed_dim = self.pretrained_embeddings.shape[1]
+
+        assert len(self.pretrained_embeddings) == len(
+            self.pretrained_concepts
+        ), "The number of embeddings does not match the number of concepts."
+
+    def _initialize_empty(self):
+        """Initialize empty attributes for when files do not exist."""
+        self.pretrained_embeddings = None
+        self.pretrained_concepts = None
+        self.concept_ids = None
+        self.concept_names = None
+        self.reverse_concept_id_map = None
+        self.embed_dim = 0
+
+    @property
+    def vocab_size(self) -> int:
+        """Return the size of the vocabulary."""
+        return len(self.pretrained_embeddings) if self.exists else 0
+
+    def is_concept_available(self, concept_id: str) -> bool:
+        """Check if a given concept ID is available."""
+        return self.exists and concept_id in self.concept_ids
+
+    def get_concept_embeddings(self, concept_id: str) -> Optional[np.ndarray]:
+        """
+        Retrieve the embedding vector for a given concept ID.
+
+        Returns None if the concept ID is not available.
+        """
+        if self.is_concept_available(concept_id):
+            return self.pretrained_embeddings[self.reverse_concept_id_map[concept_id]]
+        return None
+
+    def save(self, model_folder: str):
+        """Save the embeddings and concepts to the specified folder."""
+        if self.exists:
+            model_path = Path(model_folder)
+            np.save(
+                model_path / PRETRAINED_EMBEDDING_VECTOR_FILE_NAME,
+                self.pretrained_embeddings,
+            )
+            with open(model_path / PRETRAINED_EMBEDDING_CONCEPT_FILE_NAME, "wb") as f:
+                pickle.dump(self.pretrained_concepts, f)

cehrgpt/models/special_tokens.py

@@ -0,0 +1,30 @@
+# Special tokens for the cehr-gpt tokenizer
+START_TOKEN = "[START]"
+END_TOKEN = "[END]"
+PAD_TOKEN = "[PAD]"
+OUT_OF_VOCABULARY_TOKEN = "[OOV]"
+
+# OMOP CONCEPT IDs
+VISIT_CONCEPT_IDS = [
+    "9202",
+    "9203",
+    "581477",
+    "9201",
+    "5083",
+    "262",
+    "38004250",
+    "0",
+    "8883",
+    "38004238",
+    "38004251",
+    "38004222",
+    "38004268",
+    "38004228",
+    "32693",
+    "8971",
+    "38004269",
+    "38004193",
+    "32036",
+    "8782",
+]
+DISCHARGE_CONCEPT_IDS = []