chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,401 @@
|
||||
# Copyright (c) Facebook, Inc. and its affiliates.
|
||||
#
|
||||
# This source code is licensed under the MIT license found in the
|
||||
# LICENSE file in the root directory of this source tree.
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from fairseq import utils
|
||||
from fairseq.data import (
|
||||
Dictionary,
|
||||
IdDataset,
|
||||
ListDataset,
|
||||
NestedDictionaryDataset,
|
||||
NumelDataset,
|
||||
NumSamplesDataset,
|
||||
PadDataset,
|
||||
SortDataset,
|
||||
data_utils,
|
||||
encoders,
|
||||
)
|
||||
from fairseq.tasks import LegacyFairseqTask, register_task
|
||||
|
||||
from . import wsc_utils
|
||||
|
||||
|
||||
@register_task("wsc")
|
||||
class WSCTask(LegacyFairseqTask):
|
||||
"""Task to finetune RoBERTa for Winograd Schemas."""
|
||||
|
||||
@staticmethod
|
||||
def add_args(parser):
|
||||
"""Add task-specific arguments to the parser."""
|
||||
parser.add_argument(
|
||||
"data", metavar="DIR", help="path to data directory; we load <split>.jsonl"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--init-token",
|
||||
type=int,
|
||||
default=None,
|
||||
help="add token at the beginning of each batch item",
|
||||
)
|
||||
|
||||
def __init__(self, args, vocab):
|
||||
super().__init__(args)
|
||||
self.vocab = vocab
|
||||
self.mask = vocab.add_symbol("<mask>")
|
||||
|
||||
self.bpe = encoders.build_bpe(args)
|
||||
self.tokenizer = encoders.build_tokenizer(args)
|
||||
|
||||
# hack to handle GPT-2 BPE, which includes leading spaces
|
||||
if args.bpe == "gpt2":
|
||||
self.leading_space = True
|
||||
self.trailing_space = False
|
||||
else:
|
||||
self.leading_space = False
|
||||
self.trailing_space = True
|
||||
|
||||
@classmethod
|
||||
def load_dictionary(cls, filename):
|
||||
"""Load the dictionary from the filename
|
||||
|
||||
Args:
|
||||
filename (str): the filename
|
||||
"""
|
||||
dictionary = Dictionary.load(filename)
|
||||
dictionary.add_symbol("<mask>")
|
||||
return dictionary
|
||||
|
||||
@classmethod
|
||||
def setup_task(cls, args, **kwargs):
|
||||
assert args.criterion == "wsc", "Must set --criterion=wsc"
|
||||
|
||||
# load data and label dictionaries
|
||||
vocab = cls.load_dictionary(os.path.join(args.data, "dict.txt"))
|
||||
print("| dictionary: {} types".format(len(vocab)))
|
||||
|
||||
return cls(args, vocab)
|
||||
|
||||
def binarize(self, s: str, append_eos: bool = False):
|
||||
if self.tokenizer is not None:
|
||||
s = self.tokenizer.encode(s)
|
||||
if self.bpe is not None:
|
||||
s = self.bpe.encode(s)
|
||||
tokens = self.vocab.encode_line(
|
||||
s,
|
||||
append_eos=append_eos,
|
||||
add_if_not_exist=False,
|
||||
).long()
|
||||
if self.args.init_token is not None:
|
||||
tokens = torch.cat([tokens.new([self.args.init_token]), tokens])
|
||||
return tokens
|
||||
|
||||
def binarize_with_mask(self, txt, prefix, suffix, leading_space, trailing_space):
|
||||
toks = self.binarize(
|
||||
prefix + leading_space + txt + trailing_space + suffix,
|
||||
append_eos=True,
|
||||
)
|
||||
mask = torch.zeros_like(toks, dtype=torch.bool)
|
||||
mask_start = len(self.binarize(prefix))
|
||||
mask_size = len(self.binarize(leading_space + txt))
|
||||
mask[mask_start : mask_start + mask_size] = 1
|
||||
return toks, mask
|
||||
|
||||
def load_dataset(
|
||||
self, split, epoch=1, combine=False, data_path=None, return_only=False, **kwargs
|
||||
):
|
||||
"""Load a given dataset split.
|
||||
|
||||
Args:
|
||||
split (str): name of the split (e.g., train, valid, test)
|
||||
"""
|
||||
if data_path is None:
|
||||
data_path = os.path.join(self.args.data, split + ".jsonl")
|
||||
if not os.path.exists(data_path):
|
||||
raise FileNotFoundError("Cannot find data: {}".format(data_path))
|
||||
|
||||
query_tokens = []
|
||||
query_masks = []
|
||||
query_lengths = []
|
||||
candidate_tokens = []
|
||||
candidate_masks = []
|
||||
candidate_lengths = []
|
||||
labels = []
|
||||
|
||||
for sentence, pronoun_span, query, label in wsc_utils.jsonl_iterator(data_path):
|
||||
prefix = sentence[: pronoun_span.start].text
|
||||
suffix = sentence[pronoun_span.end :].text_with_ws
|
||||
|
||||
# spaCy spans include trailing spaces, but we need to know about
|
||||
# leading spaces for the GPT-2 BPE
|
||||
leading_space = (
|
||||
" " if sentence[: pronoun_span.start].text_with_ws.endswith(" ") else ""
|
||||
)
|
||||
trailing_space = " " if pronoun_span.text_with_ws.endswith(" ") else ""
|
||||
|
||||
# get noun phrases, excluding pronouns and anything overlapping with the query
|
||||
cand_spans = wsc_utils.filter_noun_chunks(
|
||||
wsc_utils.extended_noun_chunks(sentence),
|
||||
exclude_pronouns=True,
|
||||
exclude_query=query,
|
||||
exact_match=False,
|
||||
)
|
||||
|
||||
if query is not None:
|
||||
query_toks, query_mask = self.binarize_with_mask(
|
||||
query, prefix, suffix, leading_space, trailing_space
|
||||
)
|
||||
query_len = len(query_toks)
|
||||
else:
|
||||
query_toks, query_mask, query_len = None, None, 0
|
||||
|
||||
query_tokens.append(query_toks)
|
||||
query_masks.append(query_mask)
|
||||
query_lengths.append(query_len)
|
||||
|
||||
cand_toks, cand_masks = [], []
|
||||
for cand_span in cand_spans:
|
||||
toks, mask = self.binarize_with_mask(
|
||||
cand_span.text,
|
||||
prefix,
|
||||
suffix,
|
||||
leading_space,
|
||||
trailing_space,
|
||||
)
|
||||
cand_toks.append(toks)
|
||||
cand_masks.append(mask)
|
||||
|
||||
# collate candidates
|
||||
cand_toks = data_utils.collate_tokens(cand_toks, pad_idx=self.vocab.pad())
|
||||
cand_masks = data_utils.collate_tokens(cand_masks, pad_idx=0)
|
||||
assert cand_toks.size() == cand_masks.size()
|
||||
|
||||
candidate_tokens.append(cand_toks)
|
||||
candidate_masks.append(cand_masks)
|
||||
candidate_lengths.append(cand_toks.size(1))
|
||||
|
||||
labels.append(label)
|
||||
|
||||
query_lengths = np.array(query_lengths)
|
||||
query_tokens = ListDataset(query_tokens, query_lengths)
|
||||
query_masks = ListDataset(query_masks, query_lengths)
|
||||
|
||||
candidate_lengths = np.array(candidate_lengths)
|
||||
candidate_tokens = ListDataset(candidate_tokens, candidate_lengths)
|
||||
candidate_masks = ListDataset(candidate_masks, candidate_lengths)
|
||||
|
||||
labels = ListDataset(labels, [1] * len(labels))
|
||||
|
||||
dataset = {
|
||||
"id": IdDataset(),
|
||||
"query_tokens": query_tokens,
|
||||
"query_masks": query_masks,
|
||||
"candidate_tokens": candidate_tokens,
|
||||
"candidate_masks": candidate_masks,
|
||||
"labels": labels,
|
||||
"nsentences": NumSamplesDataset(),
|
||||
"ntokens": NumelDataset(query_tokens, reduce=True),
|
||||
}
|
||||
|
||||
nested_dataset = NestedDictionaryDataset(
|
||||
dataset,
|
||||
sizes=[query_lengths],
|
||||
)
|
||||
|
||||
with data_utils.numpy_seed(self.args.seed):
|
||||
shuffle = np.random.permutation(len(query_tokens))
|
||||
dataset = SortDataset(
|
||||
nested_dataset,
|
||||
# shuffle
|
||||
sort_order=[shuffle],
|
||||
)
|
||||
|
||||
if return_only:
|
||||
return dataset
|
||||
|
||||
self.datasets[split] = dataset
|
||||
return self.datasets[split]
|
||||
|
||||
def build_dataset_for_inference(self, sample_json):
|
||||
with tempfile.NamedTemporaryFile(buffering=0) as h:
|
||||
h.write((json.dumps(sample_json) + "\n").encode("utf-8"))
|
||||
dataset = self.load_dataset(
|
||||
"disambiguate_pronoun",
|
||||
data_path=h.name,
|
||||
return_only=True,
|
||||
)
|
||||
return dataset
|
||||
|
||||
def disambiguate_pronoun(self, model, sentence, use_cuda=False):
|
||||
sample_json = wsc_utils.convert_sentence_to_json(sentence)
|
||||
dataset = self.build_dataset_for_inference(sample_json)
|
||||
sample = dataset.collater([dataset[0]])
|
||||
if use_cuda:
|
||||
sample = utils.move_to_cuda(sample)
|
||||
|
||||
def get_masked_input(tokens, mask):
|
||||
masked_tokens = tokens.clone()
|
||||
masked_tokens[mask.bool()] = self.mask
|
||||
return masked_tokens
|
||||
|
||||
def get_lprobs(tokens, mask):
|
||||
logits, _ = model(src_tokens=get_masked_input(tokens, mask))
|
||||
lprobs = F.log_softmax(logits, dim=-1, dtype=torch.float)
|
||||
scores = lprobs.gather(2, tokens.unsqueeze(-1)).squeeze(-1)
|
||||
mask = mask.type_as(scores)
|
||||
scores = (scores * mask).sum(dim=-1) / mask.sum(dim=-1)
|
||||
return scores
|
||||
|
||||
cand_lprobs = get_lprobs(
|
||||
sample["candidate_tokens"][0],
|
||||
sample["candidate_masks"][0],
|
||||
)
|
||||
if sample["query_tokens"][0] is not None:
|
||||
query_lprobs = get_lprobs(
|
||||
sample["query_tokens"][0].unsqueeze(0),
|
||||
sample["query_masks"][0].unsqueeze(0),
|
||||
)
|
||||
return (query_lprobs >= cand_lprobs).all().item() == 1
|
||||
else:
|
||||
best_idx = cand_lprobs.argmax().item()
|
||||
full_cand = sample["candidate_tokens"][0][best_idx]
|
||||
mask = sample["candidate_masks"][0][best_idx]
|
||||
toks = full_cand[mask.bool()]
|
||||
return self.bpe.decode(self.source_dictionary.string(toks)).strip()
|
||||
|
||||
@property
|
||||
def source_dictionary(self):
|
||||
return self.vocab
|
||||
|
||||
@property
|
||||
def target_dictionary(self):
|
||||
return self.vocab
|
||||
|
||||
|
||||
@register_task("winogrande")
|
||||
class WinograndeTask(WSCTask):
|
||||
"""
|
||||
Task for WinoGrande dataset. Efficient implementation for Winograd schema
|
||||
tasks with exactly two candidates, one of which is correct.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def setup_task(cls, args, **kwargs):
|
||||
assert args.criterion == "winogrande", "Must set --criterion=winogrande"
|
||||
|
||||
# load data and label dictionaries
|
||||
vocab = cls.load_dictionary(os.path.join(args.data, "dict.txt"))
|
||||
print("| dictionary: {} types".format(len(vocab)))
|
||||
|
||||
return cls(args, vocab)
|
||||
|
||||
def load_dataset(
|
||||
self, split, epoch=1, combine=False, data_path=None, return_only=False, **kwargs
|
||||
):
|
||||
"""Load a given dataset split.
|
||||
|
||||
Args:
|
||||
split (str): name of the split (e.g., train, valid, test)
|
||||
"""
|
||||
if data_path is None:
|
||||
data_path = os.path.join(self.args.data, split + ".jsonl")
|
||||
if not os.path.exists(data_path):
|
||||
raise FileNotFoundError("Cannot find data: {}".format(data_path))
|
||||
|
||||
query_tokens = []
|
||||
query_masks = []
|
||||
query_lengths = []
|
||||
candidate_tokens = []
|
||||
candidate_masks = []
|
||||
candidate_lengths = []
|
||||
|
||||
itr = wsc_utils.winogrande_jsonl_iterator(data_path, eval=(split == "test"))
|
||||
|
||||
for sample in itr:
|
||||
sentence, pronoun_span, query, cand_text = sample
|
||||
prefix = sentence[: pronoun_span[0]].rstrip()
|
||||
suffix = sentence[pronoun_span[1] :]
|
||||
|
||||
leading_space = " " if sentence[: pronoun_span[0]].endswith(" ") else ""
|
||||
trailing_space = ""
|
||||
|
||||
if query is not None:
|
||||
query_toks, query_mask = self.binarize_with_mask(
|
||||
query,
|
||||
prefix,
|
||||
suffix,
|
||||
leading_space,
|
||||
trailing_space,
|
||||
)
|
||||
query_len = len(query_toks)
|
||||
else:
|
||||
query_toks, query_mask, query_len = None, None, 0
|
||||
|
||||
query_tokens.append(query_toks)
|
||||
query_masks.append(query_mask)
|
||||
query_lengths.append(query_len)
|
||||
|
||||
cand_toks, cand_mask = self.binarize_with_mask(
|
||||
cand_text,
|
||||
prefix,
|
||||
suffix,
|
||||
leading_space,
|
||||
trailing_space,
|
||||
)
|
||||
|
||||
candidate_tokens.append(cand_toks)
|
||||
candidate_masks.append(cand_mask)
|
||||
candidate_lengths.append(cand_toks.size(0))
|
||||
|
||||
query_lengths = np.array(query_lengths)
|
||||
|
||||
def get_pad_dataset_fn(tokens, length, pad_idx):
|
||||
return PadDataset(
|
||||
ListDataset(tokens, length),
|
||||
pad_idx=pad_idx,
|
||||
left_pad=False,
|
||||
)
|
||||
|
||||
query_tokens = get_pad_dataset_fn(query_tokens, query_lengths, self.vocab.pad())
|
||||
query_masks = get_pad_dataset_fn(query_masks, query_lengths, 0)
|
||||
|
||||
candidate_lengths = np.array(candidate_lengths)
|
||||
candidate_tokens = get_pad_dataset_fn(
|
||||
candidate_tokens, candidate_lengths, self.vocab.pad()
|
||||
)
|
||||
candidate_masks = get_pad_dataset_fn(candidate_masks, candidate_lengths, 0)
|
||||
|
||||
dataset = {
|
||||
"id": IdDataset(),
|
||||
"query_tokens": query_tokens,
|
||||
"query_masks": query_masks,
|
||||
"candidate_tokens": candidate_tokens,
|
||||
"candidate_masks": candidate_masks,
|
||||
"nsentences": NumSamplesDataset(),
|
||||
"ntokens": NumelDataset(query_tokens, reduce=True),
|
||||
}
|
||||
|
||||
nested_dataset = NestedDictionaryDataset(
|
||||
dataset,
|
||||
sizes=[query_lengths],
|
||||
)
|
||||
|
||||
with data_utils.numpy_seed(self.args.seed):
|
||||
shuffle = np.random.permutation(len(query_tokens))
|
||||
dataset = SortDataset(
|
||||
nested_dataset,
|
||||
# shuffle
|
||||
sort_order=[shuffle],
|
||||
)
|
||||
|
||||
if return_only:
|
||||
return dataset
|
||||
|
||||
self.datasets[split] = dataset
|
||||
return self.datasets[split]
|
||||
Reference in New Issue
Block a user