<commit_before>from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torchtext.prototype.models import (
    T5_BASE_GENERATION,
    T5_SMALL_GENERATION,
    T5_LARGE_GENERATION,
    T5_3B_GENERATION,
    T5_11B_GENERATION,
    T5Conf,
    T5Transform,
    T5Bundle,
)


BUNDLERS = {
    "base": T5_BASE_GENERATION,
    "small": T5_SMALL_GENERATION,
    "large": T5_LARGE_GENERATION,
    "3b": T5_3B_GENERATION,
    "11b": T5_11B_GENERATION,
}


class T5Wrapper(nn.Module):
    def __init__(
        self,
        configuration: Optional[str] = None,
        checkpoint: Optional[Union[str, Dict[str, torch.Tensor]]] = None,
        t5_config: Optional[T5Conf] = None,
        transform: Optional[T5Transform] = None,
        freeze_model: bool = False,
        strict: bool = False,
        dl_kwargs: Dict[str, Any] = None,
        device: Optional[torch.device] = None,
    ) -> None:
        """
        Args:
            configuration (str or None): The model configuration. Only support 'base', 'small', 'large', '3b', and '11b' . Must be `None` if checkpoint is not `None`. (Default: `None`)
            checkpoint (str, Dict[str, torch.Tensor], or None): Path to or actual model state_dict. state_dict can have partial weights i.e only for encoder. Must be `None` if configuration is not `None`.(Default: ``None``)
            t5_config (T5Conf or None): An instance of T5Conf that defined the model configuration (i.e. number of layer, attention heads, etc). Must be provided if configuration is `None`. (Default: `None`)
            transform (T5Transfrom or None): An instance of T5Transform that defines the text processing pipeline. Must be provided if configuration is `None`. (Default: `None`)
            freeze_model (bool): Indicates whether to freeze the model weights. (Default: `False`)
            strict (bool): Passed to :func: `torch.nn.Module.load_state_dict` method. (Default: `False`)
            dl_kwargs (dictionary of keyword arguments): Passed to :func:`torch.hub.load_state_dict_from_url`. (Default: `None`)
        """
        super().__init__()

        if configuration is None:
            assert checkpoint is not None, "Must provide a checkpoint if configuration is None"
            assert t5_config is not None, "Must provide t5_config if using checkpoint"
            assert isinstance(t5_config, T5Conf), "t5_config must have type torchtext.prototype.models.T5Conf"
            assert not t5_config.encoder_only, "t5_config.encoder_only must be False"
            assert t5_config.linear_head, "t5_config.linear_head must be True"
            assert transform is not None, "Must provide transform if using checkpoint"
            assert isinstance(transform, T5Transform), "transform must have type torchtext.prototype.models.T5Transform"<commit_msg>Make comment paths dynamic (#1894)<commit_after>from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torchtext.prototype.models import (
    T5_BASE_GENERATION,
    T5_SMALL_GENERATION,
    T5_LARGE_GENERATION,
    T5_3B_GENERATION,
    T5_11B_GENERATION,
    T5Conf,
    T5Transform,
    T5Bundle,
)


BUNDLERS = {
    "base": T5_BASE_GENERATION,
    "small": T5_SMALL_GENERATION,
    "large": T5_LARGE_GENERATION,
    "3b": T5_3B_GENERATION,
    "11b": T5_11B_GENERATION,
}


class T5Wrapper(nn.Module):
    def __init__(
        self,
        configuration: Optional[str] = None,
        checkpoint: Optional[Union[str, Dict[str, torch.Tensor]]] = None,
        t5_config: Optional[T5Conf] = None,
        transform: Optional[T5Transform] = None,
        freeze_model: bool = False,
        strict: bool = False,
        dl_kwargs: Dict[str, Any] = None,
        device: Optional[torch.device] = None,
    ) -> None:
        """
        Args:
            configuration (str or None): The model configuration. Only support 'base', 'small', 'large', '3b', and '11b' . Must be `None` if checkpoint is not `None`. (Default: `None`)
            checkpoint (str, Dict[str, torch.Tensor], or None): Path to or actual model state_dict. state_dict can have partial weights i.e only for encoder. Must be `None` if configuration is not `None`.(Default: ``None``)
            t5_config (T5Conf or None): An instance of T5Conf that defined the model configuration (i.e. number of layer, attention heads, etc). Must be provided if configuration is `None`. (Default: `None`)
            transform (T5Transfrom or None): An instance of T5Transform that defines the text processing pipeline. Must be provided if configuration is `None`. (Default: `None`)
            freeze_model (bool): Indicates whether to freeze the model weights. (Default: `False`)
            strict (bool): Passed to :func: `torch.nn.Module.load_state_dict` method. (Default: `False`)
            dl_kwargs (dictionary of keyword arguments): Passed to :func:`torch.hub.load_state_dict_from_url`. (Default: `None`)
        """
        super().__init__()

        if configuration is None:
            assert checkpoint is not None, "Must provide a checkpoint if configuration is None"
            assert t5_config is not None, "Must provide t5_config if using checkpoint"
            assert isinstance(t5_config, T5Conf), f"t5_config must have type {T5Conf.__module__}"
            assert not t5_config.encoder_only, "t5_config.encoder_only must be False"
            assert t5_config.linear_head, "t5_config.linear_head must be True"
            assert transform is not None, "Must provide transform if using checkpoint"
            assert isinstance(transform, T5Transform), f"transform must have type {T5Transform.__module__}"
{'fraction_non_alphanumeric': 0.09518051476939023, 'fraction_numerical': 0.016583174987044393, 'mean_word_length': 3.1926140477914555, 'pattern_counts': {'":': 10, '<': 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': '5616298', 'n_tokens_mistral': 1723, 'n_tokens_neox': 1698, 'n_words': 573}