from torch.utils.data import DataLoader, WeightedRandomSampler
from transformers import Trainer as HFTrainer

class WeightedTrainer(HFTrainer):
    # --- Weighted training loader ---
    def get_train_dataloader(self):
        if self.train_dataset is None:
            raise ValueError("Trainer: training requires a train_dataset.")
        
        weights = self.train_dataset.get_sample_weights()
        #print(f"Sample weights: {max(weights)} {min(weights)}")
        sampler = WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)

        return DataLoader(
            self.train_dataset,
            batch_size=self.args.per_device_train_batch_size,  # per-device batch size
            sampler=sampler,
            num_workers=self.args.dataloader_num_workers if hasattr(self.args, "dataloader_num_workers") else 0,
            pin_memory=True,
        )

    def get_eval_dataloader(self, eval_dataset=None):
        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
        if eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")

        return DataLoader(
            eval_dataset,
            batch_size=self.args.per_device_eval_batch_size,  # validation batch size
            shuffle=False,  # do not shuffle evaluation
            num_workers=self.args.dataloader_num_workers if hasattr(self.args, "dataloader_num_workers") else 0,
            pin_memory=True,
        )
