Repository: facebookresearch/ParlAI

Issue:
FairSeq grade candidates
Fairseq can currently generate and evaluate a single label, but it ignores `label_candidates` for when there's multiple labels.

The change belongs here:
https://github.com/facebookresearch/ParlAI/blob/5994d437089892ec626fcd510b6fefd2869d8a80/parlai/agents/fairseq/fairseq.py#L348-L349


Discussion:

Resolving patch:
diff --git a/parlai/agents/fairseq/fairseq.py b/parlai/agents/fairseq/fairseq.py
--- a/parlai/agents/fairseq/fairseq.py
+++ b/parlai/agents/fairseq/fairseq.py
@@ -5,6 +5,7 @@
 # of patent rights can be found in the PATENTS file in the same directory.
 
 from parlai.core.dict import DictionaryAgent
+from parlai.core.utils import argsort, padded_tensor
 
 try:
     from fairseq import models, optim, criterions
@@ -15,6 +16,7 @@
     )
 from fairseq import trainer, fp16_trainer
 from fairseq.sequence_generator import SequenceGenerator
+from fairseq.sequence_scorer import SequenceScorer
 from fairseq import options
 from fairseq.tasks.fairseq_task import FairseqTask
 from fairseq.utils import convert_padding_direction
@@ -319,7 +321,7 @@ def __init__(self, opt, shared=None):
             # actually construct the model and generator
             self.model = self.build_model()
 
-            # Construct the generator
+            # Construct the generator and scorer
             self.generator = SequenceGenerator(
                 [self.model],
                 tgt_dict=self.dict,
@@ -332,6 +334,8 @@ def __init__(self, opt, shared=None):
                 sampling_topk=self.args.sampling_topk,
                 sampling_temperature=self.args.sampling_temperature,
             )
+            self.scorer = SequenceScorer([self.model], self.dict)
+
             # set up the grader and the trainer
             self.criterion = criterions.build_criterion(self.args, self.task)
 
@@ -464,12 +468,53 @@ def eval_step(self, batch):
         if batch.label_vec is not None:
             # Interactive mode won't have a gold label
             self.trainer.valid_step(samples)
-        # Grade each of the candidate sequences
-        # TODO: grade everything in observations[i]['label_candidates']
 
+        # Output placeholders
+        reranked_cands = None
+        generated_output = None
+
+        # Grade each of the candidate sequences
+        if batch.candidate_vecs is not None:
+            bsz = len(batch.text_vec)
+            reranked_cands = []
+            # score the candidates for each item in the batch separately, so that
+            # we can support variable number of candidates
+            for i in range(bsz):
+                cands = batch.candidate_vecs[i]
+                if not cands:
+                    reranked_cands.append(None)
+                    continue
+                ncand = len(cands)
+                # repeat the input many times
+                xs = batch.text_vec[i].unsqueeze(0).expand(ncand, -1)
+                # some models crash if there's leading padding on every example
+                xs = xs[:, :batch.text_lengths[i]]
+                # and appropriately pack the outputs
+                ys, _ = padded_tensor(cands, self.NULL_IDX, self.use_cuda)
+                s = self._make_sample(xs, ys)
+                # perform the actual grading, extract the scores
+                scored = list(self.scorer.score_batched_itr([s], cuda=self.use_cuda))
+                scores = [s[3][0]['score'].item() for s in scored]
+                # intentional hanging comma here; argsort returns a list
+                ranked, = argsort(scores, batch.candidates[i], descending=True)
+                reranked_cands.append(ranked)
+
+        # Next generate freely to create our response
         if not self.args.skip_generation:
-            # Next generate freely to create our response
-            return Output(self._generate(samples), None)
+            generated_output = self._generate(samples)
+        elif reranked_cands:
+            # we're skiping generation, but we're also grading candidates
+            # so output the highest ranked candidate
+            # In the case of zero candidates, we don't have something to rank,
+            # so we may need to pass on that None
+            generated_output = [
+                ranked and ranked[0] or None for ranked in reranked_cands
+            ]
+        else:
+            # no output at all
+            pass
+
+        return Output(generated_output, reranked_cands)
 
     def _generate(self, samples):
         src_tokens = samples["net_input"]["src_tokens"]
@@ -554,6 +599,7 @@ def _make_sample(self, xs, ys):
         # TODO: should the right/left padding thing be in torch agent?
         repadded = convert_padding_direction(xs, self.dict.pad(), right_to_left=True)
         sample = {}
+        sample["id"] = torch.arange(len(xs) - 1)
         sample["net_input"] = {
             "src_tokens": repadded,
             "src_lengths": self._seq_length(xs),
diff --git a/parlai/core/torch_agent.py b/parlai/core/torch_agent.py
--- a/parlai/core/torch_agent.py
+++ b/parlai/core/torch_agent.py
@@ -6,7 +6,7 @@
 
 from parlai.core.agents import Agent
 from parlai.core.dict import DictionaryAgent
-from parlai.core.utils import set_namedtuple_defaults
+from parlai.core.utils import set_namedtuple_defaults, argsort, padded_tensor
 
 try:
     import torch
@@ -352,21 +352,13 @@ def batchify(self, obs_batch, sort=False,
         # TEXT
         xs, x_lens = None, None
         if any('text_vec' in ex for ex in exs):
-            x_text = [ex.get('text_vec', self.EMPTY) for ex in exs]
-            x_lens = [x.shape[0] for x in x_text]
-
+            _xs = [ex.get('text_vec', self.EMPTY) for ex in exs]
+            xs, x_lens = padded_tensor(_xs)
             if sort:
                 sort = False  # now we won't sort on labels
-                ind_sorted = sorted(range(len(x_lens)),
-                                    key=lambda k: -x_lens[k])
-                exs = [exs[k] for k in ind_sorted]
-                valid_inds = [valid_inds[k] for k in ind_sorted]
-                x_text = [x_text[k] for k in ind_sorted]
-                x_lens = [x_lens[k] for k in ind_sorted]
-
-            xs = torch.LongTensor(len(exs), max(x_lens)).fill_(self.NULL_IDX)
-            for i, ex in enumerate(x_text):
-                xs[i, :ex.shape[0]] = ex
+                xs, x_lens, valid_inds, exs = argsort(
+                    x_lens, xs, x_lens, valid_inds, exs, descending=True
+                )
             if self.use_cuda:
                 xs = xs.cuda()
 
@@ -384,14 +376,11 @@ def batchify(self, obs_batch, sort=False,
             y_lens = [y.shape[0] for y in label_vecs]
 
             if sort and xs is None:
-                # always sort on xs if we have them, not ys
-                ind_sorted = sorted(range(len(y_lens)),
-                                    key=lambda k: -y_lens[k])
-                exs = [exs[k] for k in ind_sorted]
-                valid_inds = [valid_inds[k] for k in ind_sorted]
-                label_vecs = [label_vecs[k] for k in ind_sorted]
-                labels = [labels[k] for k in ind_sorted]
-                y_lens = [y_lens[k] for k in ind_sorted]
+                ys, y_lens = padded_tensor(label_vecs)
+                exs, valid_inds, label_vecs, labels, y_lens = argsort(
+                    y_lens, exs, valid_inds, label_vecs, labels, y_lens,
+                    descending=True
+                )
 
             ys = torch.LongTensor(len(exs), max(y_lens)).fill_(self.NULL_IDX)
             for i, y in enumerate(label_vecs):
@@ -461,10 +450,10 @@ def match_batch(self, batch_reply, valid_inds, output=None):
     def _add_person_tokens(self, text, token, add_after_newln=False):
         if add_after_newln:
             split = text.split('\n')
-            split[-1] = token + split[-1]
+            split[-1] = token + ' ' + split[-1]
             return '\n'.join(split)
         else:
-            return token + text
+            return token + ' ' +text
 
     def get_dialog_history(self, observation, reply=None,
                            add_person_tokens=False, add_p1_after_newln=False):
@@ -490,15 +479,14 @@ def get_dialog_history(self, observation, reply=None,
         if reply is not None:
             if add_person_tokens:
                 # add person2 token to reply
-                reply = self._add_person_tokens(reply, self.P2_TOKEN + ' ')
+                reply = self._add_person_tokens(reply, self.P2_TOKEN)
             # add reply to history
             self.history.append(reply)
 
         if 'text' in obs:
             if add_person_tokens:
                 # add person1 token to text
-                obs['text'] = self._add_person_tokens(obs['text'],
-                                                      self.P1_TOKEN + ' ',
+                obs['text'] = self._add_person_tokens(obs['text'], self.P1_TOKEN,
                                                       add_p1_after_newln)
             # add text to history
             self.history.append(obs['text'])
diff --git a/parlai/core/utils.py b/parlai/core/utils.py
--- a/parlai/core/utils.py
+++ b/parlai/core/utils.py
@@ -10,6 +10,14 @@
 import random
 import time
 
+# some of the utility methods are helpful for Torch
+try:
+    import torch
+    __TORCH_AVAILABLE = True
+except ImportError:
+    __TORCH_AVAILABLE = False
+
+
 DISPLAY_MESSAGE_DEFAULT_FIELDS = {
     'episode_done',
     'id',
@@ -872,3 +880,70 @@ def set_namedtuple_defaults(namedtuple, default=None):
     """
     namedtuple.__new__.__defaults__ = (default,) * len(namedtuple._fields)
     return namedtuple
+
+
+def padded_tensor(items, pad_idx=0, use_cuda=False, left_padded=False):
+    """Create a right-padded matrix from an uneven list of lists.
+
+    Returns (padded, lengths), where padded is the padded matrix, and lengths
+    is a list containing the lengths of each row.
+
+    Matrix is right-padded (filled to the right) by default, but can be
+    left padded if the flag is set to True.
+
+    Matrix can also be placed on cuda automatically.
+
+    :param list[iter[int]] items: List of items
+    :param bool sort: If True, orders by the length
+    :param int pad_idx: the value to use for padding
+    :param bool use_cuda: if true, places `padded` on GPU
+    :param bool left_padded:
+
+    :return: (padded, lengths) tuple
+    :rtype: (Tensor[int64], list[int])
+    """
+    # hard fail if we don't have torch
+    if not __TORCH_AVAILABLE:
+        raise ImportError(
+            "Cannot use padded_tensor without torch; go to http://pytorch.org"
+        )
+
+    # number of items
+    n = len(items)
+    # length of each item
+    lens = [len(item) for item in items]
+    # max in time dimension
+    t = max(lens)
+    output = torch.LongTensor(n, t).fill_(pad_idx)
+    for i in range(len(items)):
+        if left_padded:
+            # place at end
+            output[i, t - lens[i]:] = torch.LongTensor(items[i])
+        else:
+            # place at beginning
+            output[i, :lens[i]] = torch.LongTensor(items[i])
+    if use_cuda:
+        output = output.cuda()
+    return output, lens
+
+
+def argsort(keys, *lists, descending=False):
+    """Reorder each list in lists by the (descending) sorted order of keys.
+
+    :param iter keys: Keys to order by
+    :param list[list] lists: Lists to reordered by keys's order.
+                             Correctly handles lists and 1-D tensors.
+    :param bool descending: Use descending order if true
+    :return: The reordered items
+    """
+    ind_sorted = sorted(range(len(keys)), key=lambda k: keys[k])
+    if descending:
+        ind_sorted = list(reversed(ind_sorted))
+    output = []
+    for lst in lists:
+        # watch out in case we don't have torch installed
+        if __TORCH_AVAILABLE and isinstance(lst, torch.Tensor):
+            output.append(lst[ind_sorted])
+        else:
+            output.append([lst[i] for i in ind_sorted])
+    return output


Test patch:
diff --git a/tests/run_tests_short.sh b/tests/run_tests_short.sh
--- a/tests/run_tests_short.sh
+++ b/tests/run_tests_short.sh
@@ -9,6 +9,7 @@
 set -e # stop if any tests fail
 for test in $(ls); do
     if [ ${test: -3} == ".py" ] && [ $test != "test_downloads.py" ] && [ $test != "test_mlb_vqa.py" ]; then
+        echo "Running $test"
         python3 $test;
     fi
 done
diff --git a/tests/test_torch_agent.py b/tests/test_torch_agent.py
--- a/tests/test_torch_agent.py
+++ b/tests/test_torch_agent.py
@@ -520,12 +520,12 @@ def test__add_person_tokens(self):
             "Attack ships on fire off the shoulder of Orion.\n"
             "I watched C-beams glitter in the dark near the Tannhauser gate.\n"
             "All those moments will be lost in time, like tears in rain.")
-        prefix = 'PRE '
+        prefix = 'PRE'
         out = agent._add_person_tokens(text, prefix, add_after_newln=False)
-        self.assertEqual(out, prefix + text)
+        self.assertEqual(out, prefix + ' ' + text)
         out = agent._add_person_tokens(text, prefix, add_after_newln=True)
         idx = text.rfind('\n') + 1
-        self.assertEqual(out, text[:idx] + prefix + text[idx:])
+        self.assertEqual(out, text[:idx] + prefix + ' ' + text[idx:])
 
     def test_get_dialog_history(self):
         """Test different dialog history settings."""
diff --git a/tests/test_utils.py b/tests/test_utils.py
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -4,13 +4,18 @@
 # LICENSE file in the root directory of this source tree. An additional grant
 # of patent rights can be found in the PATENTS file in the same directory.
 
-from parlai.core.utils import Timer, round_sigfigs, set_namedtuple_defaults
+from parlai.core.utils import Timer
+from parlai.core.utils import round_sigfigs
+from parlai.core.utils import set_namedtuple_defaults
+from parlai.core.utils import padded_tensor
+from parlai.core.utils import argsort
 import time
 import unittest
+import torch
+import numpy as np
 
 
 class TestUtils(unittest.TestCase):
-
     def test_round_sigfigs(self):
         x = 0
         y = 0
@@ -89,3 +94,30 @@ def test_setnamedtupledefaults(self):
         assert nt.a is 1
         assert nt.b is 1
         assert nt.c is 1
+
+    def test_padded_tensor(self):
+        # list of lists
+        lol = [[1, 2], [3, 4, 5]]
+        output, lens = padded_tensor(lol)
+        assert np.all(output.numpy() == np.array([[1, 2, 0], [3, 4, 5]]))
+        assert lens == [2, 3]
+        output, _ = padded_tensor(lol, left_padded=True)
+        assert np.all(output.numpy() == np.array([[0, 1, 2], [3, 4, 5]]))
+        output, _ = padded_tensor(lol, pad_idx=99)
+        assert np.all(output.numpy() == np.array([[1, 2, 99], [3, 4, 5]]))
+
+    def test_argsort(self):
+        keys = [5, 4, 3, 2, 1]
+        items = ["five", "four", "three", "two", "one"]
+        items2 = ["e", "d", "c", "b", "a"]
+        torch_keys = torch.LongTensor(keys)
+        assert argsort(keys, items, items2) == [
+            list(reversed(items)), list(reversed(items2))
+        ]
+        assert argsort(keys, items, items2, descending=True) == [items, items2]
+
+        assert np.all(argsort(torch_keys, torch_keys)[0].numpy() == np.arange(1, 6))
+
+
+if __name__ == '__main__':
+    unittest.main()