chendl/compositional_test

1import itertools2import json3import linecache4import os5import pickle6import re7import socket8import string9from collections import Counter10from logging import getLogger11from pathlib import Path12from typing import Callable, Dict, Iterable, List13 14import git15import torch16from torch.utils.data import Dataset17 18from transformers import BartTokenizer, RagTokenizer, T5Tokenizer19 20 21def encode_line(tokenizer, line, max_length, padding_side, pad_to_max_length=True, return_tensors="pt"):22    extra_kw = {"add_prefix_space": True} if isinstance(tokenizer, BartTokenizer) and not line.startswith(" ") else {}23    tokenizer.padding_side = padding_side24    return tokenizer(25        [line],26        max_length=max_length,27        padding="max_length" if pad_to_max_length else None,28        truncation=True,29        return_tensors=return_tensors,30        add_special_tokens=True,31        **extra_kw,32    )33 34 35def trim_batch(36    input_ids,37    pad_token_id,38    attention_mask=None,39):40    """Remove columns that are populated exclusively by pad_token_id"""41    keep_column_mask = input_ids.ne(pad_token_id).any(dim=0)42    if attention_mask is None:43        return input_ids[:, keep_column_mask]44    else:45        return (input_ids[:, keep_column_mask], attention_mask[:, keep_column_mask])46 47 48class Seq2SeqDataset(Dataset):49    def __init__(50        self,51        tokenizer,52        data_dir,53        max_source_length,54        max_target_length,55        type_path="train",56        n_obs=None,57        src_lang=None,58        tgt_lang=None,59        prefix="",60    ):61        super().__init__()62        self.src_file = Path(data_dir).joinpath(type_path + ".source")63        self.tgt_file = Path(data_dir).joinpath(type_path + ".target")64        self.src_lens = self.get_char_lens(self.src_file)65        self.max_source_length = max_source_length66        self.max_target_length = max_target_length67        assert min(self.src_lens) > 0, f"found empty line in {self.src_file}"68        self.tokenizer = tokenizer69        self.prefix = prefix70        if n_obs is not None:71            self.src_lens = self.src_lens[:n_obs]72        self.src_lang = src_lang73        self.tgt_lang = tgt_lang74 75    def __len__(self):76        return len(self.src_lens)77 78    def __getitem__(self, index) -> Dict[str, torch.Tensor]:79        index = index + 1  # linecache starts at 180        source_line = self.prefix + linecache.getline(str(self.src_file), index).rstrip("\n")81        tgt_line = linecache.getline(str(self.tgt_file), index).rstrip("\n")82        assert source_line, f"empty source line for index {index}"83        assert tgt_line, f"empty tgt line for index {index}"84 85        # Need to add eos token manually for T586        if isinstance(self.tokenizer, T5Tokenizer):87            source_line += self.tokenizer.eos_token88            tgt_line += self.tokenizer.eos_token89 90        # Pad source and target to the right91        source_tokenizer = (92            self.tokenizer.question_encoder if isinstance(self.tokenizer, RagTokenizer) else self.tokenizer93        )94        target_tokenizer = self.tokenizer.generator if isinstance(self.tokenizer, RagTokenizer) else self.tokenizer95 96        source_inputs = encode_line(source_tokenizer, source_line, self.max_source_length, "right")97        target_inputs = encode_line(target_tokenizer, tgt_line, self.max_target_length, "right")98 99        source_ids = source_inputs["input_ids"].squeeze()100        target_ids = target_inputs["input_ids"].squeeze()101        src_mask = source_inputs["attention_mask"].squeeze()102        return {103            "input_ids": source_ids,104            "attention_mask": src_mask,105            "decoder_input_ids": target_ids,106        }107 108    @staticmethod109    def get_char_lens(data_file):110        return [len(x) for x in Path(data_file).open().readlines()]111 112    def collate_fn(self, batch) -> Dict[str, torch.Tensor]:113        input_ids = torch.stack([x["input_ids"] for x in batch])114        masks = torch.stack([x["attention_mask"] for x in batch])115        target_ids = torch.stack([x["decoder_input_ids"] for x in batch])116        tgt_pad_token_id = (117            self.tokenizer.generator.pad_token_id118            if isinstance(self.tokenizer, RagTokenizer)119            else self.tokenizer.pad_token_id120        )121        src_pad_token_id = (122            self.tokenizer.question_encoder.pad_token_id123            if isinstance(self.tokenizer, RagTokenizer)124            else self.tokenizer.pad_token_id125        )126        y = trim_batch(target_ids, tgt_pad_token_id)127        source_ids, source_mask = trim_batch(input_ids, src_pad_token_id, attention_mask=masks)128        batch = {129            "input_ids": source_ids,130            "attention_mask": source_mask,131            "decoder_input_ids": y,132        }133        return batch134 135 136logger = getLogger(__name__)137 138 139def flatten_list(summary_ids: List[List]):140    return list(itertools.chain.from_iterable(summary_ids))141 142 143def save_git_info(folder_path: str) -> None:144    """Save git information to output_dir/git_log.json"""145    repo_infos = get_git_info()146    save_json(repo_infos, os.path.join(folder_path, "git_log.json"))147 148 149def save_json(content, path, indent=4, **json_dump_kwargs):150    with open(path, "w") as f:151        json.dump(content, f, indent=indent, **json_dump_kwargs)152 153 154def load_json(path):155    with open(path) as f:156        return json.load(f)157 158 159def get_git_info():160    repo = git.Repo(search_parent_directories=True)161    repo_infos = {162        "repo_id": str(repo),163        "repo_sha": str(repo.head.object.hexsha),164        "repo_branch": str(repo.active_branch),165        "hostname": str(socket.gethostname()),166    }167    return repo_infos168 169 170def lmap(f: Callable, x: Iterable) -> List:171    """list(map(f, x))"""172    return list(map(f, x))173 174 175def pickle_save(obj, path):176    """pickle.dump(obj, path)"""177    with open(path, "wb") as f:178        return pickle.dump(obj, f)179 180 181def normalize_answer(s):182    """Lower text and remove punctuation, articles and extra whitespace."""183 184    def remove_articles(text):185        return re.sub(r"\b(a|an|the)\b", " ", text)186 187    def white_space_fix(text):188        return " ".join(text.split())189 190    def remove_punc(text):191        exclude = set(string.punctuation)192        return "".join(ch for ch in text if ch not in exclude)193 194    def lower(text):195        return text.lower()196 197    return white_space_fix(remove_articles(remove_punc(lower(s))))198 199 200def f1_score(prediction, ground_truth):201    prediction_tokens = normalize_answer(prediction).split()202    ground_truth_tokens = normalize_answer(ground_truth).split()203    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)204    num_same = sum(common.values())205    if num_same == 0:206        return 0207    precision = 1.0 * num_same / len(prediction_tokens)208    recall = 1.0 * num_same / len(ground_truth_tokens)209    f1 = (2 * precision * recall) / (precision + recall)210    return f1211 212 213def exact_match_score(prediction, ground_truth):214    return normalize_answer(prediction) == normalize_answer(ground_truth)215 216 217def calculate_exact_match(output_lns: List[str], reference_lns: List[str]) -> Dict:218    assert len(output_lns) == len(reference_lns)219    em = 0220    for hypo, pred in zip(output_lns, reference_lns):221        em += exact_match_score(hypo, pred)222    if len(output_lns) > 0:223        em /= len(output_lns)224    return {"em": em}225 226 227def is_rag_model(model_prefix):228    return model_prefix.startswith("rag")229 230 231def set_extra_model_params(extra_params, hparams, config):232    equivalent_param = {p: p for p in extra_params}233    # T5 models don't have `dropout` param, they have `dropout_rate` instead234    equivalent_param["dropout"] = "dropout_rate"235    for p in extra_params:236        if getattr(hparams, p, None):237            if not hasattr(config, p) and not hasattr(config, equivalent_param[p]):238                logger.info("config doesn't have a `{}` attribute".format(p))239                delattr(hparams, p)240                continue241            set_p = p if hasattr(config, p) else equivalent_param[p]242            setattr(config, set_p, getattr(hparams, p))243            delattr(hparams, p)244    return hparams, config245