Remove extra info token pathway

This commit is contained in:
2026-07-07 16:57:49 +08:00
parent 6dfeb5a696
commit a0379daf29
13 changed files with 18 additions and 1390 deletions

View File

@@ -1,6 +1,6 @@
from __future__ import annotations
from typing import Any, Dict, Iterable, List
from typing import Any, Dict, List
import numpy as np
import torch
@@ -25,7 +25,6 @@ class AllFutureSequenceEvalDataset:
labels_file: str,
min_history_events: int = 1,
min_future_events: int = 1,
extra_info_types: Iterable[int] | None = None,
) -> None:
base = AllFutureHealthDataset(
data_prefix=data_prefix,
@@ -33,18 +32,12 @@ class AllFutureSequenceEvalDataset:
split="train",
min_history_events=min_history_events,
min_future_events=min_future_events,
extra_info_types=extra_info_types,
)
self.base = base
self.label_code_to_id = base.label_code_to_id
self.label_id_to_code = base.label_id_to_code
self.vocab_size = base.vocab_size
self.n_types = base.n_types
self.n_cont_types = base.n_cont_types
self.n_categories = base.n_categories
self.cont_type_ids = base.cont_type_ids
self.extra_info_types = base.extra_info_types
self.samples: List[Dict[str, Any]] = []
for patient in base.patients:
@@ -62,10 +55,6 @@ class AllFutureSequenceEvalDataset:
"target_time_seq": times[1:],
"readout_mask": np.ones(input_len, dtype=bool),
"sex": int(patient["sex"]),
"other_type": np.asarray(patient["other_type"], dtype=np.int64),
"other_value": np.asarray(patient["other_value"], dtype=np.float32),
"other_value_kind": np.asarray(patient["other_value_kind"], dtype=np.int64),
"other_time": np.asarray(patient["other_time"], dtype=np.float32),
}
)
@@ -81,10 +70,6 @@ class AllFutureSequenceEvalDataset:
"target_time_seq": torch.from_numpy(s["target_time_seq"]).float(),
"readout_mask": torch.from_numpy(s["readout_mask"]).bool(),
"sex": torch.tensor(s["sex"], dtype=torch.long),
"other_type": torch.from_numpy(s["other_type"]).long(),
"other_value": torch.from_numpy(s["other_value"]).float(),
"other_value_kind": torch.from_numpy(s["other_value_kind"]).long(),
"other_time": torch.from_numpy(s["other_time"]).float(),
}
@@ -97,7 +82,6 @@ def load_sequence_eval_dataset(
include_no_event_in_uts_target: bool,
min_history_events: int,
min_future_events: int,
extra_info_types: Iterable[int] | None,
):
mode = str(model_target_mode).lower()
if mode == "next_token":
@@ -106,7 +90,6 @@ def load_sequence_eval_dataset(
labels_file=labels_file,
no_event_interval_years=no_event_interval_years,
include_no_event_in_uts_target=include_no_event_in_uts_target,
extra_info_types=extra_info_types,
)
if mode == "all_future":
return AllFutureSequenceEvalDataset(
@@ -114,7 +97,6 @@ def load_sequence_eval_dataset(
labels_file=labels_file,
min_history_events=min_history_events,
min_future_events=min_future_events,
extra_info_types=extra_info_types,
)
raise ValueError(f"Unknown model_target_mode: {model_target_mode!r}")
@@ -135,19 +117,6 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str,
readout_mask = pad_sequence(
[s["readout_mask"] for s in batch], batch_first=True, padding_value=False
)
other_type = pad_sequence(
[s["other_type"] for s in batch], batch_first=True, padding_value=0
)
other_value = pad_sequence(
[s["other_value"] for s in batch], batch_first=True, padding_value=0.0
)
other_value_kind = pad_sequence(
[s["other_value_kind"] for s in batch], batch_first=True, padding_value=0
)
other_time = pad_sequence(
[s["other_time"] for s in batch], batch_first=True, padding_value=0.0
)
return {
"event_seq": event_seq,
"time_seq": time_seq,
@@ -156,8 +125,4 @@ def sequence_eval_collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str,
"target_time_seq": target_time_seq,
"readout_mask": readout_mask,
"sex": torch.stack([s["sex"] for s in batch]),
"other_type": other_type,
"other_value": other_value,
"other_value_kind": other_value_kind,
"other_time": other_time,
}