tComp/tcomp_utils.py at master · Sharif-SLPL/tComp · 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
import torch
from dataclasses import dataclass
from typing import List, Optional, Tuple, Union


@dataclass
class tcompConfig():
    include_biases: Optional[bool] = True
    bias_decomp_type: Optional[str] = "absdot"  # "absdot": Based on the absolute value of dot products | "norm": Based on the norm of the attribution vectors | "equal": equal decomposition | "abssim": Based on the absolute value of cosine similarites | "cls": add to cls token
    include_bias_token: Optional[bool] = False  # Adds an extra input token as a bias in the attribution vectors
    # If the bias_decomp_type is None and include_bias_token is True then the final token in the input tokens of the attr. vectors will be the summation of the biases
    # Otherwise the bias token will be decomposed with the specified decomp type

    include_LN1: Optional[bool] = True

    include_FFN: Optional[bool] = True
    FFN_approx_type: Optional[str] = "GeLU_ZO"  # "GeLU_LA": GeLU-based linear approximation | "ReLU": Using ReLU as an approximation | "GeLU_ZO": Zero-origin slope approximation
    FFN_fast_mode: Optional[bool] = False

    include_LN2: Optional[bool] = True

    aggregation: Optional[str] = None  # None: No aggregation | vector: Vector-based aggregation | rollout: Norm-based rollout aggregation

    include_classifier_w_pooler: Optional[bool] = True
    tanh_approx_type: Optional[str] = "ZO"  # "ZO": Zero-origin slope approximation | "LA": Linear approximation

    output_enc_all_layers: Optional[bool] = False  # True: Output all layers | False: Output only last layer
    output_enc_attention: Optional[str] = None  # None | norm | vector | both
    output_enc_res1: Optional[str] = None  # None | norm | vector | both
    output_enc_LN1: Optional[str] = None  # None | norm | vector | both
    output_enc_FFN: Optional[str] = None  # None | norm | vector | both
    output_enc_res2: Optional[str] = None  # None | norm | vector | both
    output_enc_encoder: Optional[str] = None  # None | norm | vector | both
    # output_aggregated: Optional[str] = None  # None | norm | vector | both
    # output_pooler: Optional[str] = None  # None | norm | vector | both
    # output_classifier: Optional[bool] = True

    output_dec_all_layers: Optional[bool] = False  # True: Output all layers | False: Output only last layer
    output_dec_self_attention: Optional[str] = None  # None | norm | vector | both
    output_dec_res1: Optional[str] = None  # None | norm | vector | both
    output_dec_LN1: Optional[str] = None  # None | norm | vector | both
    output_dec_cross_attention: Optional[str] = None  # None | norm | vector | both
    output_dec_res2: Optional[str] = None  # None | norm | vector | both
    output_dec_LN2: Optional[str] = None  # None | norm | vector | both
    output_dec_FFN: Optional[str] = None  # None | norm | vector | both
    output_dec_res3: Optional[str] = None  # None | norm | vector | both
    output_dec_decoder: Optional[str] = None  # None | norm | vector | both


@dataclass
class tcompOutput_Enc():
    attention: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    res1: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    LN1: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    FFN: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    res2: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    encoder: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    # aggregated: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    # pooler: Optional[Union[Tuple[torch.Tensor], torch.Tensor]] = None
    # classifier: Optional[torch.Tensor] = None


@dataclass
class tcompOutput_Dec():
    self_attention: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    res1: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    LN1: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    cross_attention: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    res2: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    LN2: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    FFN: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    res3: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None
    decoder: Optional[Union[Tuple[torch.Tensor, Tuple[torch.Tensor]], Tuple[torch.Tensor], torch.Tensor]] = None