h2o-llmstudio

1
import logging
2
from typing import Any, Dict
3

4
from torch import nn
5
from transformers import AutoModelForCausalLM
6

7
from llm_studio.src.utils.data_utils import batch_padding
8
from llm_studio.src.utils.modeling_utils import create_nlp_backbone, prepare_lora
9

10
logger = logging.getLogger(__name__)
11

12

13
class Model(nn.Module):
14
    """
15
    Model for causal language modeling problem type.
16
    """
17

18
    def __init__(self, cfg: Any):
19
        """
20
        Args:
21
            cfg: config with all the hyperparameters
22
        """
23

24
        super(Model, self).__init__()
25

26
        self.cfg = cfg
27
        self.backbone, self.backbone_config = create_nlp_backbone(
28
            cfg, model_class=AutoModelForCausalLM
29
        )
30

31
        if cfg.training.lora:
32
            self.backbone = prepare_lora(cfg, self.backbone)
33

34
        self.classification_head = nn.Linear(
35
            self.backbone_config.vocab_size, cfg.dataset.num_classes, bias=False
36
        )
37

38
        self.loss_fn = self.cfg.training.loss_class.get(
39
            self.cfg.training.loss_function
40
        )(self.cfg)
41

42
    def forward(
43
        self,
44
        batch: Dict,
45
        padding: bool = True,
46
    ) -> Dict:
47
        # disable cache if gradient checkpointing is enabled
48
        if self.cfg.architecture.gradient_checkpointing:
49
            self.backbone.config.use_cache = False
50

51
        outputs: Dict = {}
52
        mask_key = "prompt_attention_mask"
53
        pad_keys = [
54
            "prompt_input_ids",
55
            "prompt_attention_mask",
56
            "special_tokens_mask",
57
            "labels",
58
        ]
59

60
        if padding:
61
            batch = batch_padding(
62
                self.cfg,
63
                batch,
64
                self.training,
65
                mask_key=mask_key,
66
                pad_keys=pad_keys,
67
                padding_side=self.cfg.tokenizer._padding_side,
68
            )
69

70
        output = self.backbone(
71
            input_ids=batch["prompt_input_ids"],
72
            attention_mask=batch["prompt_attention_mask"],
73
        )
74

75
        output.logits = self.classification_head(output[0][:, -1].float())
76

77
        if "labels" in batch:
78
            loss = self.loss_fn(
79
                output.logits, batch["class_label"].unsqueeze(1).float()
80
            )
81
            outputs["loss"] = loss
82

83
        outputs["logits"] = output.logits
84

85
        # enable cache again if gradient checkpointing is enabled
86
        if self.cfg.architecture.gradient_checkpointing:
87
            self.backbone.config.use_cache = True
88

89
        return outputs
90