Function evaluate_model

train.py:24–48

Evaluate the model on the eval dataset for each flag and return losses.

(model, eval_dataloader, flags)

    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels, "flag": flag}

def evaluate_model(model, eval_dataloader, flags):
    """Evaluate the model on the eval dataset for each flag and return losses."""
    model.eval()
    eval_losses = {flag: 0.0 for flag in flags}
    num_batches = len(eval_dataloader)

    with torch.no_grad():
        for flag in flags:
            total_loss = 0.0
            for batch in eval_dataloader:
                input_ids = batch["input_ids"].to(model.device)
                attention_mask = batch["attention_mask"].to(model.device)
                labels = batch["labels"].to(model.device)

                # Configure the subnetwork for the flag
                model.configure_subnetwork(flag)

                # Forward pass
                outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
                total_loss += outputs.loss.item()

            eval_losses[flag] = total_loss / num_batches

    model.train()
    return eval_losses

if __name__ == "__main__":
    # Load tokenizer