|
|
|
|
|
|
|
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
|
|
|
