| """ |
| Fine-tuning data pipeline for JointFusionModel. |
| |
| Prepares datasets that yield {input_ids, attention_mask, tabular_features, labels} |
| for supervised fine-tuning of the joint Transformer + DCNv2(PLR) fusion model. |
| |
| Unlike pre-training (which packs sequences for 100% token utilization), fine-tuning |
| uses per-user sequences padded to a fixed length — each sample represents one user |
| with one label. |
| """ |
|
|
| import logging |
| from typing import Any, Dict, Sequence |
|
|
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset as TorchDataset |
|
|
| from ..tokenizers.domain_tokenizer import DomainTokenizerBuilder |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class DomainFinetuneDataset(TorchDataset): |
| """Dataset for fine-tuning JointFusionModel on labeled user data. |
| |
| Each sample represents one user: |
| - input_ids: tokenized transaction sequence (padded/truncated to max_length) |
| - attention_mask: 1 for real tokens, 0 for padding |
| - tabular_features: numerical feature vector for the user |
| - labels: target label (float for binary, int for multiclass) |
| """ |
|
|
| def __init__(self, user_sequences, tabular_features, labels, builder, hf_tokenizer, max_length=512): |
| assert len(user_sequences) == len(tabular_features) == len(labels), ( |
| f"Length mismatch: sequences={len(user_sequences)}, " |
| f"tabular={len(tabular_features)}, labels={len(labels)}" |
| ) |
| self.user_sequences = user_sequences |
| self.tabular_features = np.asarray(tabular_features, dtype=np.float32) |
| self.labels = np.asarray(labels) |
| self.builder = builder |
| self.hf_tokenizer = hf_tokenizer |
| self.max_length = max_length |
| if hf_tokenizer.pad_token_id is None: |
| raise ValueError("Tokenizer must have pad_token set.") |
|
|
| def __len__(self): |
| return len(self.user_sequences) |
|
|
| def __getitem__(self, idx): |
| events = self.user_sequences[idx] |
| token_strings = self.builder.tokenize_sequence(events, add_bos=True, add_eos=True) |
| encoding = self.hf_tokenizer( |
| " ".join(token_strings), max_length=self.max_length, |
| truncation=True, padding="max_length", add_special_tokens=False, return_tensors="pt", |
| ) |
| return { |
| "input_ids": encoding["input_ids"].squeeze(0), |
| "attention_mask": encoding["attention_mask"].squeeze(0), |
| "tabular_features": torch.tensor(self.tabular_features[idx], dtype=torch.float32), |
| "labels": torch.tensor(self.labels[idx], dtype=torch.float32), |
| } |
|
|
| def get_stats(self): |
| return { |
| "n_samples": len(self), "max_length": self.max_length, |
| "n_tabular_features": self.tabular_features.shape[1], |
| "label_distribution": { |
| "mean": float(self.labels.mean()), "std": float(self.labels.std()), |
| "min": float(self.labels.min()), "max": float(self.labels.max()), |
| }, |
| } |
|
|
|
|
| def prepare_finetune_dataset(user_sequences, tabular_features, labels, builder, hf_tokenizer, max_length=512): |
| """Convenience function to create a fine-tuning dataset.""" |
| ds = DomainFinetuneDataset(user_sequences, tabular_features, labels, builder, hf_tokenizer, max_length) |
| logger.info(f"Fine-tune dataset: {len(ds)} samples, max_length={max_length}, " |
| f"tabular_features={tabular_features.shape[1]}, label_mean={labels.mean():.3f}") |
| return ds |
|
|
