<commit_before>from typing import Callable, List

import torch.nn as nn
from torch.utils.data import DataLoader


class Task(object):
    """A task for use in an MMTL MetalModel

    Args:
        name: The name of the task
        TODO: replace this with a more fully-featured path through the network
        input_module: The input module
        head_module: The task head module
        data: A list of DataLoaders (instances and labels) to feed through the network.
            The list contains [train, dev, test].
        scorers: A list of Scorers that return metrics_dict objects.
    """

    def __init__(
        self,
        name: str,
        input_module: nn.Module,
        head_module: nn.Module,
        data_loaders: List[DataLoader],
        scorers: List[Callable] = None,
    ) -> None:
        if len(data_loaders) != 3:
            msg = "Arg data_loaders must be a list of length 3 [train, valid, test]"
            raise Exception(msg)
        self.name = name
        self.input_module = input_module
        self.head_module = head_module
        self.data_loaders = data_loaders
        self.scorers = scorers<commit_msg>Update Task signature<commit_after>from typing import Callable, List

import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader


class Task(object):
    """A task for use in an MMTL MetalModel

    Args:
        name: The name of the task
        TODO: replace this with a more fully-featured path through the network
        input_module: The input module
        head_module: The task head module
        data: A list of DataLoaders (instances and labels) to feed through the network.
            The list contains [train, dev, test].
        scorers: A list of Scorers that return metrics_dict objects.
    """

    def __init__(
        self,
        name: str,
        data_loaders: List[DataLoader],
        input_module: nn.Module,
        head_module: nn.Module,
        scorers: List[Callable] = None,
        loss_hat_func: Callable = F.cross_entropy,
        probs_hat_func: Callable = F.softmax,
    ) -> None:
        if len(data_loaders) != 3:
            msg = "Arg data_loaders must be a list of length 3 [train, valid, test]"
            raise Exception(msg)
        self.name = name
        self.data_loaders = data_loaders
        self.input_module = input_module
        self.head_module = head_module
        self.scorers = scorers
        self.loss_hat_func = loss_hat_func
        self.probs_hat_func = probs_hat_func
{'fraction_non_alphanumeric': 0.06880189798339265, 'fraction_numerical': 0.0015816528272044287, 'mean_word_length': 2.324572930354796, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '722058', 'n_tokens_mistral': 740, 'n_tokens_neox': 718, 'n_words': 257}