189 lines
7.2 KiB
Python
189 lines
7.2 KiB
Python
import torch
|
|
import numpy as np
|
|
import random
|
|
from collections import defaultdict
|
|
import json
|
|
from models import TimeAwareGPT2
|
|
|
|
|
|
class PatientEventDataset(torch.utils.data.Dataset):
|
|
"""
|
|
A PyTorch Dataset for handling temporal sequences of patient events.
|
|
|
|
This class processes a raw NumPy array of patient records, groups them by
|
|
patient ID, and prepares them for training by imputing gaps, padding, or
|
|
truncating sequences to a fixed length.
|
|
"""
|
|
|
|
def __init__(self, data: np.ndarray, block_length: int):
|
|
"""
|
|
Initializes the dataset by pre-processing the patient event data.
|
|
|
|
Args:
|
|
data (np.ndarray): A NumPy array of shape (N, 3) with dtype=np.uint32.
|
|
The columns represent (patient_id, time_in_days, event_code).
|
|
block_length (int): The fixed length for the output sequences.
|
|
"""
|
|
self.block_length = block_length
|
|
|
|
# Group (time_in_days, event_code) pairs by patient_id.
|
|
# This pre-processing step allows for efficient lookups in __getitem__.
|
|
patient_events = defaultdict(list)
|
|
for patient_id, time, event in data:
|
|
patient_events[patient_id].append((time, event))
|
|
|
|
# Store a list of unique patient_ids to map indices to patients.
|
|
self.patient_ids = list(patient_events.keys())
|
|
self.patient_events = dict(patient_events)
|
|
|
|
def __len__(self) -> int:
|
|
"""
|
|
Returns the total number of unique patients in the dataset.
|
|
"""
|
|
return len(self.patient_ids)
|
|
|
|
def __getitem__(self, idx):
|
|
"""
|
|
Retrieves, processes, and returns a single patient's event sequence,
|
|
or a list of sequences if a slice is provided.
|
|
|
|
Args:
|
|
idx (int or slice): The index or slice of the patient(s) to retrieve.
|
|
|
|
Returns:
|
|
If idx is an int, a tuple of two torch.long tensors:
|
|
(event_sequence, time_sequence), both of shape (block_length,).
|
|
If idx is a slice, a list of such tuples.
|
|
"""
|
|
if isinstance(idx, slice):
|
|
return [self[i] for i in range(*idx.indices(len(self)))]
|
|
|
|
# 1. Retrieve and Sort
|
|
patient_id = self.patient_ids[idx]
|
|
records = sorted(self.patient_events[patient_id], key=lambda x: x[0])
|
|
|
|
# 2. Impute "No Event" Gaps
|
|
imputed_sequence = []
|
|
if not records:
|
|
# Handle cases with no records for a patient if necessary, though
|
|
# the constructor logic would typically prevent this.
|
|
pass
|
|
else:
|
|
imputed_sequence.append(records[0])
|
|
for i in range(len(records) - 1):
|
|
prev_time, _ = records[i]
|
|
next_time, _ = records[i+1]
|
|
time_gap = next_time - prev_time
|
|
|
|
# If the gap is 5 years (1826 days) or more, insert "no event" records.
|
|
if time_gap >= 1826:
|
|
num_no_event_intervals = time_gap // 1826
|
|
for j in range(1, num_no_event_intervals + 1):
|
|
no_event_time = prev_time + j * 1826
|
|
imputed_sequence.append((no_event_time, 1)) # event_code=1 for "no event"
|
|
|
|
imputed_sequence.append(records[i+1])
|
|
|
|
# 3. Adjust Sequence Length
|
|
seq_len = len(imputed_sequence)
|
|
|
|
if seq_len > self.block_length:
|
|
# If longer, randomly select a contiguous sub-sequence.
|
|
start_index = random.randint(0, seq_len - self.block_length)
|
|
final_sequence = imputed_sequence[start_index : start_index + self.block_length]
|
|
elif seq_len < self.block_length:
|
|
# If shorter, pad the sequence at the end.
|
|
padding_needed = self.block_length - seq_len
|
|
# Use event_code=0 and time_in_days=36525 for padding.
|
|
padding = [(36525, 0)] * padding_needed
|
|
final_sequence = imputed_sequence + padding
|
|
else:
|
|
# If equal, use the sequence as is.
|
|
final_sequence = imputed_sequence
|
|
|
|
# 4. Return Tensors
|
|
# Separate the sequence into event codes and time, then convert to tensors.
|
|
event_codes = [item[1] for item in final_sequence]
|
|
time_stamps = [item[0] for item in final_sequence]
|
|
|
|
event_tensor = torch.tensor(event_codes, dtype=torch.long)
|
|
time_tensor = torch.tensor(time_stamps, dtype=torch.long)
|
|
|
|
return event_tensor, time_tensor
|
|
|
|
def load_model(config_path, model_path, vocab_size, device='cpu'):
|
|
"""
|
|
Loads a trained TimeAwareGPT2 model from a configuration file and a state dictionary.
|
|
|
|
Args:
|
|
config_path (str): Path to the JSON configuration file.
|
|
model_path (str): Path to the saved model state dictionary (.pt file).
|
|
vocab_size (int): The vocabulary size used during training.
|
|
device (str): The device to load the model onto ('cpu' or 'cuda').
|
|
|
|
Returns:
|
|
(TimeAwareGPT2): The loaded and initialized model.
|
|
"""
|
|
with open(config_path, 'r') as f:
|
|
config_dict = json.load(f)
|
|
|
|
print(f"Model config: {config_dict}")
|
|
|
|
# Create a config object from the dictionary
|
|
class AttrDict(dict):
|
|
def __init__(self, *args, **kwargs):
|
|
super(AttrDict, self).__init__(*args, **kwargs)
|
|
self.__dict__ = self
|
|
|
|
config = AttrDict(config_dict)
|
|
|
|
# Initialize the model with parameters from the config
|
|
model = TimeAwareGPT2(
|
|
vocab_size=vocab_size,
|
|
n_embd=config.n_embd,
|
|
n_layer=config.n_layer,
|
|
n_head=config.n_head,
|
|
pdrop=config.pdrop,
|
|
token_pdrop=config.token_pdrop
|
|
).to(device)
|
|
|
|
# Load the saved state dictionary
|
|
model.load_state_dict(torch.load(model_path, map_location=device))
|
|
|
|
# Set the model to evaluation mode
|
|
model.eval()
|
|
|
|
print(f"Model loaded from {model_path} with {model.get_num_params():.2f}M parameters.")
|
|
return model
|
|
|
|
|
|
def get_batch(dataset: PatientEventDataset, batch_slice: slice) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Retrieves a batch of data from a PatientEventDataset and prepares it for model training.
|
|
|
|
Args:
|
|
dataset (PatientEventDataset): The dataset to retrieve data from.
|
|
batch_slice (slice): The slice defining the batch of patients to retrieve.
|
|
ignore_tokens (list, optional): A list of token IDs to be ignored in the target events.
|
|
These tokens will be replaced with -100. Defaults to None.
|
|
|
|
Returns:
|
|
A tuple containing four tensors:
|
|
- input_events: (batch_size, sequence_length - 1)
|
|
- input_tims: (batch_size, sequence_length - 1)
|
|
- target_events: (batch_size, sequence_length - 1)
|
|
- target_times: (batch_size, sequence_length - 1)
|
|
"""
|
|
batch_data = dataset[batch_slice]
|
|
|
|
input_events = [item[0][:-1] for item in batch_data]
|
|
input_tims = [item[1][:-1] for item in batch_data]
|
|
target_events = [item[0][1:] for item in batch_data]
|
|
target_times = [item[1][1:] for item in batch_data]
|
|
|
|
input_events = torch.stack(input_events)
|
|
input_tims = torch.stack(input_tims)
|
|
target_events = torch.stack(target_events)
|
|
target_times = torch.stack(target_times)
|
|
|
|
return input_events, input_tims, target_events, target_times |