| |
|
| |
|
| |
|
| | from typing import Optional, Tuple
|
| | from dataclasses import dataclass
|
| | from transformers.modeling_outputs import ModelOutput
|
| | import sys, copy, math
|
| |
|
| | from .pooling import *
|
| | from .loss import *
|
| |
|
| | @dataclass
|
| | class AllOutput(ModelOutput):
|
| | losses: Optional[dict[str, dict[str, torch.FloatTensor]]] = None
|
| | outputs: Optional[dict[str, dict[str, torch.FloatTensor]]] = None
|
| | hidden_states: Optional[Tuple[torch.FloatTensor]] = None
|
| | attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| | cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| | global_attentions: Optional[Tuple[torch.FloatTensor]] = None
|
| | contacts: Optional[Tuple[torch.FloatTensor]] = None
|
| | losses_b: Optional[dict[str, dict[str, torch.FloatTensor]]] = None
|
| | outputs_b: Optional[dict[str, dict[str, torch.FloatTensor]]] = None
|
| | hidden_states_b: Optional[Tuple[torch.FloatTensor]] = None
|
| | attentions_b: Optional[Tuple[torch.FloatTensor]] = None
|
| | cross_attentions_b: Optional[Tuple[torch.FloatTensor]] = None
|
| | global_attentions_b: Optional[Tuple[torch.FloatTensor]] = None
|
| | contacts_b: Optional[Tuple[torch.FloatTensor]] = None
|
| | pair_outputs: Optional[Tuple[torch.FloatTensor]] = None
|
| | pair_losses: Optional[dict[str, dict[str, torch.FloatTensor]]] = None
|
| |
|
| |
|
| | def create_pooler(task_level_type, task_level_name, config, args):
|
| | '''
|
| | pooler building
|
| | :param task_level_type:
|
| | :param task_level_name:
|
| | :param config:
|
| | :param args:
|
| | :return:
|
| | '''
|
| | hidden_size = config.hidden_size[task_level_type][task_level_name]
|
| | pooling_type = args.pooling_type[task_level_type][task_level_name]
|
| |
|
| | if pooling_type == "max":
|
| | return GlobalMaskMaxPooling1D()
|
| | elif pooling_type == "sum":
|
| | return GlobalMaskSumPooling1D(axis=1)
|
| | elif pooling_type == "avg":
|
| | return GlobalMaskAvgPooling1D()
|
| | elif pooling_type == "attention":
|
| | return GlobalMaskContextAttentionPooling1D(embed_size=hidden_size)
|
| | elif pooling_type == "context_attention":
|
| | return GlobalMaskContextAttentionPooling1D(embed_size=hidden_size)
|
| | elif pooling_type == "weighted_attention":
|
| | return GlobalMaskWeightedAttentionPooling1D(embed_size=hidden_size)
|
| | elif pooling_type == "value_attention":
|
| | return GlobalMaskValueAttentionPooling1D(embed_size=hidden_size)
|
| | elif pooling_type == "transformer":
|
| | copy_config = copy.deepcopy(config)
|
| | copy_config.hidden_size = hidden_size
|
| | return GlobalMaskTransformerPooling1D(copy_config)
|
| | else:
|
| | return None
|
| |
|
| |
|
| | def create_output_loss_lucagplm(task_level_type, task_level_name, config):
|
| | '''not cls module'''
|
| | if not hasattr(config, "sigmoid"):
|
| | config.sigmoid = {task_level_type: {}}
|
| | elif task_level_type not in config.sigmoid:
|
| | config.sigmoid[task_level_type] = {}
|
| | config.sigmoid[task_level_type][task_level_name] = False if config.output_mode[task_level_type][task_level_name] \
|
| | in ["multi_class", "multi-class", "regression"] else True
|
| |
|
| | if task_level_name == "prot_contact":
|
| | config.sigmoid[task_level_type][task_level_name] = True
|
| | config.num_labels = config.label_size[task_level_type][task_level_name]
|
| | if task_level_type in ["token_level", "whole_level"]:
|
| | return_types = ["output", "loss"]
|
| | else:
|
| | return_types = ["dropout", "hidden_layer", "hidden_act", "classifier", "output", "loss"]
|
| | return create_loss_function(config,
|
| | task_level_type=task_level_type,
|
| | task_level_name=task_level_name,
|
| | sigmoid=config.sigmoid[task_level_type][task_level_name],
|
| | output_mode=config.output_mode[task_level_type][task_level_name],
|
| | num_labels=config.num_labels,
|
| | loss_type=config.loss_type[task_level_type][task_level_name],
|
| | ignore_index=config.ignore_index,
|
| | pair_level=True if task_level_type == "pair_level" else False,
|
| | return_types=return_types)
|
| |
|
| |
|
| | def create_output_loss(task_level_type, task_level_name, cls_module, config, args):
|
| | cls = None
|
| | if task_level_type in ["token_level", "whole_level"]:
|
| | cls = cls_module(config)
|
| | dropout, hidden_layer, hidden_act, classifier, output, loss_fct = create_output_loss_lucagplm(task_level_type, task_level_name, config, args)
|
| | return cls, dropout, hidden_layer, hidden_act, classifier, output, loss_fct
|
| |
|
