TFkit/tests/test_task_generation.py at 2c42e3da4f591c914fd459364efb923462144829 · voidful/TFkit · GitHub

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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
from types import SimpleNamespace

import torch
from torch import nn

from tfkit.task.clm.model import Model as CLMModel
from tfkit.task.seq2seq.model import Model as Seq2SeqModel


class DummyTokenizer:
    def __init__(self, vocab_size):
        self.vocab_size = vocab_size

    def __len__(self):
        return self.vocab_size

    def convert_ids_to_tokens(self, idx):
        return f"token-{idx}"


class DummyCausalPretrained(nn.Module):
    def __init__(self):
        super().__init__()
        self.config = SimpleNamespace(vocab_size=5, hidden_size=4)
        self.output_layer = nn.Linear(self.config.hidden_size, self.config.vocab_size)
        self.last_kwargs = None

    def get_output_embeddings(self):
        return self.output_layer

    def forward(self, input_ids, attention_mask=None, return_dict=True, **kwargs):
        self.last_kwargs = kwargs
        batch_size, seq_len = input_ids.shape
        logits = torch.zeros(batch_size, seq_len, self.config.vocab_size)
        outputs = {
            "logits": logits,
            "last_hidden_state": torch.zeros(batch_size, seq_len, self.config.hidden_size),
        }
        if "labels" in kwargs:
            outputs["loss"] = torch.tensor(0.0)
        return outputs


class DummyEncoderPretrained(nn.Module):
    def __init__(self):
        super().__init__()
        self.config = SimpleNamespace(vocab_size=5, hidden_size=4)
        self.last_kwargs = None

    def get_output_embeddings(self):
        return None

    def forward(self, input_ids, attention_mask=None, return_dict=True, **kwargs):
        self.last_kwargs = kwargs
        batch_size, seq_len = input_ids.shape
        hidden = torch.zeros(batch_size, seq_len, self.config.hidden_size)
        return {"last_hidden_state": hidden}


class DummySeq2SeqPretrained(nn.Module):
    def __init__(self):
        super().__init__()
        self.config = SimpleNamespace(vocab_size=3, hidden_size=4)
        self.decoder = nn.Module()
        self.output_layer = nn.Linear(self.config.hidden_size, self.config.vocab_size)

    def get_output_embeddings(self):
        return self.output_layer

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        decoder_input_ids=None,
        decoder_attention_mask=None,
        output_hidden_states=False,
        use_cache=False,
        return_dict=True,
        **kwargs,
    ):
        batch_size, seq_len = decoder_input_ids.shape
        hidden = torch.zeros(batch_size, seq_len, self.config.hidden_size)
        outputs = {
            "last_hidden_state": hidden,
            "decoder_hidden_states": (hidden,),
        }
        return outputs


def test_clm_model_uses_pretrained_head_for_loss():
    tokenizer = DummyTokenizer(vocab_size=5)
    pretrained = DummyCausalPretrained()
    model = CLMModel(tokenizer=tokenizer, pretrained=pretrained)

    batch = {
        "input": torch.zeros((1, 2), dtype=torch.long),
        "mask": torch.ones((1, 2), dtype=torch.long),
        "target": torch.tensor([[0, -1]]),
    }

    loss = model.forward(batch, eval=False)
    assert torch.is_tensor(loss)
    assert "labels" in pretrained.last_kwargs
    assert pretrained.last_kwargs["labels"].tolist() == [[0, -100]]

    eval_batch = {
        **batch,
        "start": [0],
    }
    result = model.forward(eval_batch, eval=True)
    assert isinstance(result, dict)
    assert "max_item" in result


def test_clm_model_falls_back_to_linear_head():
    tokenizer = DummyTokenizer(vocab_size=5)
    pretrained = DummyEncoderPretrained()
    model = CLMModel(tokenizer=tokenizer, pretrained=pretrained)

    batch = {
        "input": torch.zeros((1, 2), dtype=torch.long),
        "mask": torch.ones((1, 2), dtype=torch.long),
        "target": torch.tensor([[0, -1]]),
    }

    loss = model.forward(batch, eval=False)
    assert torch.is_tensor(loss)
    assert pretrained.last_kwargs == {}


def test_seq2seq_model_uses_pretrained_output_head():
    tokenizer = DummyTokenizer(vocab_size=3)
    pretrained = DummySeq2SeqPretrained()
    model = Seq2SeqModel(tokenizer=tokenizer, pretrained=pretrained)

    batch = {
        "input": torch.zeros((1, 1), dtype=torch.long),
        "prev": torch.zeros((1, 1), dtype=torch.long),
        "encoder_mask": torch.ones((1, 1), dtype=torch.long),
        "decoder_mask": torch.ones((1, 1), dtype=torch.long),
        "target": torch.zeros((1, 1), dtype=torch.long),
        "ntarget": torch.full((1, 1), -1),
    }

    loss = model.forward(batch, eval=False)
    assert torch.is_tensor(loss)
    assert model.model is pretrained.output_layer