social-data-attribution-w2/tests/unlearning/test_utils.py

# pyright: reportPrivateImportUsage=false
"""Tests for shuffle_segments and per_document_loss."""

from __future__ import annotations

import torch

from unlearning.trainer.utils import per_document_loss, shuffle_segments, shuffle_tokens


class TestShuffleSegments:
    def test_preserves_length(self):
        ids = torch.randint(10, 1000, (2, 128))
        shuffled = shuffle_segments(ids, pad_token_id=0, n_factor=10)
        assert shuffled.shape == ids.shape

    def test_preserves_token_set(self):
        ids = torch.randint(10, 500, (1, 64))
        shuffled = shuffle_segments(ids, pad_token_id=0, n_factor=5)
        assert sorted(ids[0].tolist()) == sorted(shuffled[0].tolist())

    def test_does_not_change_padding(self):
        ids = torch.tensor([[100, 200, 300, 1, 1, 1]])
        shuffled = shuffle_segments(ids, pad_token_id=1, n_factor=3)
        assert (shuffled[0, 3:] == 1).all()
        assert set(shuffled[0, :3].tolist()) == {100, 200, 300}

    def test_empty_sequence(self):
        ids = torch.tensor([[1, 1, 1]])
        shuffled = shuffle_segments(ids, pad_token_id=1, n_factor=10)
        assert (shuffled == ids).all()

    def test_larger_segments_than_shuffle_tokens(self):
        torch.manual_seed(42)
        ids = torch.randint(10, 1000, (1, 200))
        seg = shuffle_segments(ids.clone(), pad_token_id=0, n_factor=10)
        torch.manual_seed(42)
        chk = shuffle_tokens(ids.clone(), pad_token_id=0, max_chunk_len=10)
        assert not (seg == chk).all(), "segments and chunks should differ"


class TestPerDocumentLoss:
    def test_uniform_length(self):
        B, T, V = 4, 16, 100
        logits = torch.randn(B, T, V)
        labels = torch.randint(0, V, (B, T))
        loss = per_document_loss(logits, labels)
        assert loss.dim() == 0
        assert loss.item() > 0

    def test_variable_length_with_padding(self):
        B, T, V = 3, 20, 50
        logits = torch.randn(B, T, V)
        labels = torch.full((B, T), -100, dtype=torch.long)
        labels[0, :15] = torch.randint(0, V, (15,))
        labels[1, :10] = torch.randint(0, V, (10,))
        labels[2, :20] = torch.randint(0, V, (20,))
        loss = per_document_loss(logits, labels)
        assert loss.dim() == 0
        assert loss.item() > 0

    def test_all_padding_handled(self):
        B, T, V = 2, 10, 50
        logits = torch.randn(B, T, V)
        labels = torch.full((B, T), -100, dtype=torch.long)
        labels[0, :5] = torch.randint(0, V, (5,))
        loss = per_document_loss(logits, labels)
        assert not torch.isnan(loss)
        assert not torch.isinf(loss)

    def test_single_doc(self):
        B, T, V = 1, 8, 30
        logits = torch.randn(B, T, V)
        labels = torch.randint(0, V, (B, T))
        loss = per_document_loss(logits, labels)
        assert loss.dim() == 0