File size: 1,504 Bytes
46cc63a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | """Final Squeeze protocol — freeze mode and threshold grid."""
from __future__ import annotations
import numpy as np
import pytest
from src.evaluation.threshold_tuning import search_best_threshold
def test_threshold_grid_030_to_070():
y = np.array([0, 0, 1, 1, 1, 0])
probs = np.array([0.2, 0.4, 0.55, 0.62, 0.8, 0.35])
t, score = search_best_threshold(
y,
probs,
metric="f1_weighted",
min_threshold=0.30,
max_threshold=0.70,
step=0.01,
)
assert 0.30 <= t <= 0.70
assert score >= 0.0
def test_average_state_dicts():
import torch
from src.models.transformer_trainer import _average_state_dicts
a = {"w": torch.tensor([1.0, 3.0])}
b = {"w": torch.tensor([3.0, 5.0])}
avg = _average_state_dicts([a, b])
assert torch.allclose(avg["w"], torch.tensor([2.0, 4.0]))
def test_apply_model_freeze_full_mode():
pytest.importorskip("transformers")
from transformers import AutoModelForSequenceClassification
from src.models.transformer_trainer import _apply_model_freeze
model = AutoModelForSequenceClassification.from_pretrained(
"distilbert-base-uncased",
num_labels=2,
)
mode, freeze_n = _apply_model_freeze(model, {"freeze_mode": "full"})
assert mode == "full_unfreeze"
assert freeze_n == 0
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
total = sum(p.numel() for p in model.parameters())
assert trainable == total
|