<commit_before>from collections import defaultdict
from os import PathLike
from typing import (
    Any,
    Callable,
    Collection,
    Dict,
    Iterable,
    Iterator,
    List,
    Optional,
    Sequence,
    Tuple,
    Union,
)

import torch

import utils


def _apply(s: str, f: Optional[Callable[[str], str]]) -> str:
    return f(s) if f is not None else s


class Preprocessor:
    serialize_embeddings: bool = False

    def __init__(self):
        self.vocabs: Dict[str, utils.data.Vocab] = {}
        self._embeddings: Optional[List[List[float]]] = None
        self._embed_file: Optional[Union[str, bytes, PathLike]] = None

    def build_vocab(
        self,
        file: Union[str, bytes, PathLike],
        unknown: str = "<UNK>",
        preprocess: Optional[Callable[[str], str]] = str.lower,
        min_frequency: int = 2,
    ) -> None:
        word_counter: Dict[str, int] = defaultdict(int)
        postag_vocab = utils.data.Vocab(unknown)
        deprel_vocab = utils.data.Vocab()
        for tokens in utils.conll.read_conll(file):
            for token in tokens:
                word_counter[_apply(token["form"], preprocess)] += 1
                postag_vocab(token["postag"])
                deprel_vocab(token["deprel"])
        word_iter = (k for k, v in word_counter.items() if v >= min_frequency)
        postag_vocab.freeze()
        deprel_vocab.freeze()

        self.vocabs["word"] = utils.data.Vocab.fromkeys(word_iter, unknown)
        self.vocabs["word"].preprocess = preprocess
        self.vocabs["postag"] = postag_vocab
        self.vocabs["deprel"] = deprel_vocab
        if "pretrained_word" not in self.vocabs:
            self.vocabs["pretrained_word"] = utils.data.Vocab.fromkeys([], unknown)

    def load_embeddings(
        self,
        file: Union[str, bytes, PathLike],
        unknown: str = "<UNK>",
        preprocess: Optional[Callable[[str], str]] = str.lower,
    ) -> None:
        embeddings = utils.data.load_embeddings(file)
        dim = len(next(iter(embeddings.values())))
        embeddings[_apply("<ROOT>", preprocess)] = [0.0] * dim
        if unknown not in embeddings:
            embeddings[unknown] = [0.0] * dim
        self.vocabs["pretrained_word"] = utils.data.Vocab.fromkeys(embeddings.keys(), unknown)
        self.vocabs["pretrained_word"].preprocess = preprocess
        self._embeddings = list(embeddings.values())
        self._embed_file = file

    def transform(
        self, tokens: Iterable[Dict[str, Any]]
    ) -> Tuple[List[int], List[int], List[int], List[int], List[int]]:
        words, postags, heads, deprels = zip(
            *[(token["form"], token["postag"], token["head"], token["deprel"]) for token in tokens]
        )
        word_ids = [self.vocabs["word"][s] for s in words]
        pre_ids = [self.vocabs["pretrained_word"][s] for s in words]
        postag_ids = [self.vocabs["postag"][s] for s in postags]
        deprel_ids = [self.vocabs["deprel"][s] for s in deprels]
        return (word_ids, pre_ids, postag_ids, list(heads), deprel_ids)

    def __getstate__(self) -> Dict[str, Any]:
        state = self.__dict__.copy()
        if not self.serialize_embeddings:
            state["_embeddings"] = None
        return state

    @property
    def pretrained_word_embeddings(self) -> Optional[List[List[float]]]:
        if self._embeddings is None and self._embed_file is not None:
            v = self.vocabs["pretrained_word"]
            assert v.unknown_id is not None
            self.load_embeddings(self._embed_file, v.lookup(v.unknown_id), v.preprocess)
            assert len(self.vocabs["pretrained_word"]) == len(v)
        return self._embeddings


def create_dataloader(
    file: Union[str, bytes, PathLike],
    preprocessor: Preprocessor,
    device: Optional[torch.device] = None,
    **kwargs,
) -> torch.utils.data.DataLoader:
    dataset = ListDataset(map(preprocessor.transform, utils.conll.read_conll(file)))
    if kwargs.get("batch_sampler") is None:
        kwargs["batch_sampler"] = BucketSampler(
            dataset,
            key=lambda x: len(x[0]),
            batch_size=kwargs.pop("batch_size", 1),
            shuffle=kwargs.pop("shuffle", False),
            drop_last=kwargs.pop("drop_last", False),
            generator=kwargs.get("generator"),
        )
    kwargs.setdefault("collate_fn", lambda batch: collate(batch, device))
    loader = torch.utils.data.DataLoader(dataset, **kwargs)
    return loader


class ListDataset(list, torch.utils.data.Dataset):
    pass


class BucketSampler(torch.utils.data.Sampler[List[int]]):
    def __init__(<commit_msg>Cache pretrained embeddings.<commit_after>import os
from collections import defaultdict
from os import PathLike
from typing import (
    Any,
    Callable,
    Collection,
    Dict,
    Iterable,
    Iterator,
    List,
    Optional,
    Sequence,
    Tuple,
    Union,
)

import torch

import utils


def _apply(s: str, f: Optional[Callable[[str], str]]) -> str:
    return f(s) if f is not None else s


class Preprocessor:
    serialize_embeddings: bool = False

    def __init__(self):
        self.vocabs: Dict[str, utils.data.Vocab] = {}
        self._embeddings: Optional[List[List[float]]] = None
        self._embed_file: Optional[Union[str, bytes, PathLike]] = None

    def build_vocab(
        self,
        file: Union[str, bytes, PathLike],
        unknown: str = "<UNK>",
        preprocess: Optional[Callable[[str], str]] = str.lower,
        min_frequency: int = 2,
    ) -> None:
        word_counter: Dict[str, int] = defaultdict(int)
        postag_vocab = utils.data.Vocab(unknown)
        deprel_vocab = utils.data.Vocab()
        for tokens in utils.conll.read_conll(file):
            for token in tokens:
                word_counter[_apply(token["form"], preprocess)] += 1
                postag_vocab(token["postag"])
                deprel_vocab(token["deprel"])
        word_iter = (k for k, v in word_counter.items() if v >= min_frequency)
        postag_vocab.freeze()
        deprel_vocab.freeze()

        self.vocabs["word"] = utils.data.Vocab.fromkeys(word_iter, unknown)
        self.vocabs["word"].preprocess = preprocess
        self.vocabs["postag"] = postag_vocab
        self.vocabs["deprel"] = deprel_vocab
        if "pretrained_word" not in self.vocabs:
            self.vocabs["pretrained_word"] = utils.data.Vocab.fromkeys([], unknown)

    def load_embeddings(
        self,
        file: Union[str, bytes, PathLike],
        unknown: str = "<UNK>",
        preprocess: Optional[Callable[[str], str]] = str.lower,
        cache_dir: Optional[Union[str, bytes, PathLike]] = None,
    ) -> None:
        def _add_entry(token):
            if token not in vocab:
                nonlocal embeddings
                vocab.append(token)
                embeddings = torch.vstack((embeddings, torch.zeros_like(embeddings[0])))

        vocab, embeddings = _wrap_cache(_load_embeddings, file, cache_dir)
        _add_entry(_apply("<ROOT>", preprocess))
        _add_entry(unknown)
        self.vocabs["pretrained_word"] = utils.data.Vocab.fromkeys(vocab, unknown)
        self.vocabs["pretrained_word"].preprocess = preprocess
        self._embeddings = embeddings.tolist()
        self._embed_file = file

    def transform(
        self, tokens: Iterable[Dict[str, Any]]
    ) -> Tuple[List[int], List[int], List[int], List[int], List[int]]:
        words, postags, heads, deprels = zip(
            *[(token["form"], token["postag"], token["head"], token["deprel"]) for token in tokens]
        )
        word_ids = [self.vocabs["word"][s] for s in words]
        pre_ids = [self.vocabs["pretrained_word"][s] for s in words]
        postag_ids = [self.vocabs["postag"][s] for s in postags]
        deprel_ids = [self.vocabs["deprel"][s] for s in deprels]
        return (word_ids, pre_ids, postag_ids, list(heads), deprel_ids)

    def __getstate__(self) -> Dict[str, Any]:
        state = self.__dict__.copy()
        if not self.serialize_embeddings:
            state["_embeddings"] = None
        return state

    @property
    def pretrained_word_embeddings(self) -> Optional[List[List[float]]]:
        if self._embeddings is None and self._embed_file is not None:
            v = self.vocabs["pretrained_word"]
            assert v.unknown_id is not None
            self.load_embeddings(self._embed_file, v.lookup(v.unknown_id), v.preprocess)
            assert len(self.vocabs["pretrained_word"]) == len(v)
        return self._embeddings


def _load_embeddings(file):
    embeddings = utils.data.load_embeddings(file)
    return (list(embeddings.keys()), torch.tensor(list(embeddings.values())))


def _wrap_cache(load_fn, file, cache_dir=None, suffix=".cache"):
    if cache_dir is None:
        return load_fn(file)

    basename = os.path.basename(file)
    if not basename:
        raise ValueError(f"Invalid filename: '{file}'")
    cache_file = os.path.join(cache_dir, f"{basename}{suffix}")

    if os.path.exists(cache_file):
        obj = torch.load(cache_file)
    else:
        obj = load_fn(file)
        torch.save(obj, cache_file)
    return obj


def create_dataloader(
    file: Union[str, bytes, PathLike],
    preprocessor: Preprocessor,
    device: Optional[torch.device] = None,
    **kwargs,
) -> torch.utils.data.DataLoader:
    dataset = ListDataset(map(preprocessor.transform, utils.conll.read_conll(file)))
    if kwargs.get("batch_sampler") is None:
        kwargs["batch_sampler"] = BucketSampler(
            dataset,
            key=lambda x: len(x[0]),
            batch_size=kwargs.pop("batch_size", 1),
            shuffle=kwargs.pop("shuffle", False),
            drop_last=kwargs.pop("drop_last", False),
            generator=kwargs.get("generator"),
        )
    kwargs.setdefault("collate_fn", lambda batch: collate(batch, device))
    loader = torch.utils.data.DataLoader(dataset, **kwargs)
    return loader


class ListDataset(list, torch.utils.data.Dataset):
    pass


class BucketSampler(torch.utils.data.Sampler[List[int]]):
    def __init__(
{'fraction_non_alphanumeric': 0.12862823061630219, 'fraction_numerical': 0.0012922465208747515, 'mean_word_length': 2.792310591782887, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 25, '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': '1940106', 'n_tokens_mistral': 3273, 'n_tokens_neox': 3231, 'n_words': 682}