<commit_before># 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 argparse
import tempfile
import unittest

import tests.utils as test_utils
import torch
from fairseq import search
from fairseq.data.dictionary import Dictionary

from fairseq.models.transformer import TransformerModel
from fairseq.sequence_generator import SequenceGenerator, EnsembleModel
from fairseq.tasks.fairseq_task import FairseqTask


DEFAULT_TEST_VOCAB_SIZE = 100


class DummyTask(FairseqTask):
    def __init__(self, args):
        super().__init__(args)
        self.dictionary = get_dummy_dictionary()
        if getattr(self.args, "ctc", False):
            self.dictionary.add_symbol("<ctc_blank>")
        self.src_dict = self.dictionary
        self.tgt_dict = self.dictionary

    @property
    def source_dictionary(self):
        return self.src_dict

    @property
    def target_dictionary(self):
        return self.dictionary


def get_dummy_dictionary(vocab_size=DEFAULT_TEST_VOCAB_SIZE):
    dummy_dict = Dictionary()
    # add dummy symbol to satisfy vocab size
    for id, _ in enumerate(range(vocab_size)):
        dummy_dict.add_symbol("{}".format(id), 1000)
    return dummy_dict


def get_dummy_task_and_parser():
    """
    to build a fariseq model, we need some dummy parse and task. This function
    is used to create dummy task and parser to faciliate model/criterion test

    Note: we use FbSpeechRecognitionTask as the dummy task. You may want
    to use other task by providing another function
    """
    parser = argparse.ArgumentParser(
        description="test_dummy_s2s_task", argument_default=argparse.SUPPRESS
    )
    DummyTask.add_args(parser)
    args = parser.parse_args([])
    task = DummyTask.setup_task(args)
    return task, parser


class TestJitSequenceGeneratorBase(unittest.TestCase):
    def setUp(self):
        self.task, self.parser = get_dummy_task_and_parser()
        eos = self.task.tgt_dict.eos()
        src_tokens = torch.randint(3, 50, (2, 10)).long()
        src_tokens = torch.cat((src_tokens, torch.LongTensor([[eos], [eos]])), -1)
        src_lengths = torch.LongTensor([2, 10])
        self.sample = {
            "net_input": {"src_tokens": src_tokens, "src_lengths": src_lengths}
        }
        TransformerModel.add_args(self.parser)
        args = self.parser.parse_args([])
        args.encoder_layers = 2
        args.decoder_layers = 1
        self.transformer_model = TransformerModel.build_model(args, self.task)

    def assertOutputEqual(self, hypo, pos_probs):
        pos_scores = torch.FloatTensor(pos_probs).log()
        self.assertTensorSizeEqual(hypo["positional_scores"], pos_scores)
        self.assertTensorSizeEqual(pos_scores.numel(), hypo["tokens"].numel())

    def assertTensorSizeEqual(self, t1, t2):
        self.assertEqual(t1.size(), t2.size(), "size mismatch")

    def assertAlmostEqual(self, t1, t2):
        self.assertEqual(t1.size(), t2.size(), "size mismatch")
        self.assertLess((t1 - t2).abs().max(), 1e-4)

    def assertTensorEqual(self, t1, t2):
        self.assertEqual(t1.size(), t2.size(), "size mismatch")
        self.assertEqual(t1.ne(t2).long().sum(), 0)

    def assertHypoEqual(self, h1, h2):
        "Check two hypos are equal"
        self.assertTensorEqual(h1["tokens"], h2["tokens"])
        self.assertAlmostEqual(h1["positional_scores"], h2["positional_scores"])
        self.assertLess(abs(h1["score"] - h2["score"]), 1e-6)
        self.assertAlmostEqual(h1["attention"], h2["attention"])

    def _test_save_and_load(self, scripted_module):
        with tempfile.NamedTemporaryFile() as f:
            scripted_module.save(f.name)
            torch.jit.load(f.name)


class TestJitSequeneceGenerator(TestJitSequenceGeneratorBase):

    @unittest.skipIf(
        torch.__version__ < "1.5.0", "Targeting OSS scriptability for the 1.5 release"
    )
    def test_export_transformer(self):
        model = self.transformer_model
        torch.jit.script(model)

    @unittest.skipIf(
        torch.__version__ < "1.5.0", "Targeting OSS scriptability for the 1.5 release"
    )
    def test_ensemble_sequence_generator(self):
        model = self.transformer_model
        generator = SequenceGenerator(
            [model], self.task.tgt_dict, beam_size=2, no_repeat_ngram_size=2
        )
        scripted_model = torch.jit.script(generator)
        self._test_save_and_load(scripted_model)

    @unittest.skipIf(
        torch.__version__ < "1.5.0", "Targeting OSS scriptability for the 1.5 release"
    )
    def test_quantized_ensemble_sequence_generator(self):
        model = torch.quantization.quantize_dynamic(
            self.transformer_model, {torch.nn.Linear}, dtype=torch.qint8, inplace=True
        )
        generator = SequenceGenerator([model], self.task.tgt_dict, beam_size=2)
        scripted_model = torch.jit.script(generator)
        self._test_save_and_load(scripted_model)


class TestJitEnsemble(TestJitSequenceGeneratorBase):

    @unittest.skipIf(
        torch.__version__ < "1.5.0", "Targeting OSS scriptability for the 1.5 release"
    )
    def test_export_ensemble_model(self):
        model = self.transformer_model
        ensemble_models = EnsembleModel([model])
        torch.jit.script(ensemble_models)


class TestExportSearch(unittest.TestCase):
    def setUp(self):
        task, _ = get_dummy_task_and_parser()
        self.tgt_dict = task.tgt_dict
        self.min_top1_prob = 0.4

    def test_export_diverse_bs(self):<commit_msg>Modify gated unit tests to fix Fairseq OSS (#2059)<commit_after># 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 argparse
import tempfile
import unittest

import tests.utils as test_utils
import torch
from fairseq import search
from fairseq.data.dictionary import Dictionary

from fairseq.models.transformer import TransformerModel
from fairseq.sequence_generator import SequenceGenerator, EnsembleModel
from fairseq.tasks.fairseq_task import FairseqTask


DEFAULT_TEST_VOCAB_SIZE = 100


class DummyTask(FairseqTask):
    def __init__(self, args):
        super().__init__(args)
        self.dictionary = get_dummy_dictionary()
        if getattr(self.args, "ctc", False):
            self.dictionary.add_symbol("<ctc_blank>")
        self.src_dict = self.dictionary
        self.tgt_dict = self.dictionary

    @property
    def source_dictionary(self):
        return self.src_dict

    @property
    def target_dictionary(self):
        return self.dictionary


def get_dummy_dictionary(vocab_size=DEFAULT_TEST_VOCAB_SIZE):
    dummy_dict = Dictionary()
    # add dummy symbol to satisfy vocab size
    for id, _ in enumerate(range(vocab_size)):
        dummy_dict.add_symbol("{}".format(id), 1000)
    return dummy_dict


def get_dummy_task_and_parser():
    """
    to build a fariseq model, we need some dummy parse and task. This function
    is used to create dummy task and parser to faciliate model/criterion test

    Note: we use FbSpeechRecognitionTask as the dummy task. You may want
    to use other task by providing another function
    """
    parser = argparse.ArgumentParser(
        description="test_dummy_s2s_task", argument_default=argparse.SUPPRESS
    )
    DummyTask.add_args(parser)
    args = parser.parse_args([])
    task = DummyTask.setup_task(args)
    return task, parser


class TestJitSequenceGeneratorBase(unittest.TestCase):
    def setUp(self):
        self.task, self.parser = get_dummy_task_and_parser()
        eos = self.task.tgt_dict.eos()
        src_tokens = torch.randint(3, 50, (2, 10)).long()
        src_tokens = torch.cat((src_tokens, torch.LongTensor([[eos], [eos]])), -1)
        src_lengths = torch.LongTensor([2, 10])
        self.sample = {
            "net_input": {"src_tokens": src_tokens, "src_lengths": src_lengths}
        }
        TransformerModel.add_args(self.parser)
        args = self.parser.parse_args([])
        args.encoder_layers = 2
        args.decoder_layers = 1
        self.transformer_model = TransformerModel.build_model(args, self.task)

    def assertOutputEqual(self, hypo, pos_probs):
        pos_scores = torch.FloatTensor(pos_probs).log()
        self.assertTensorSizeEqual(hypo["positional_scores"], pos_scores)
        self.assertTensorSizeEqual(pos_scores.numel(), hypo["tokens"].numel())

    def assertTensorSizeEqual(self, t1, t2):
        self.assertEqual(t1.size(), t2.size(), "size mismatch")

    def assertAlmostEqual(self, t1, t2):
        self.assertEqual(t1.size(), t2.size(), "size mismatch")
        self.assertLess((t1 - t2).abs().max(), 1e-4)

    def assertTensorEqual(self, t1, t2):
        self.assertEqual(t1.size(), t2.size(), "size mismatch")
        self.assertEqual(t1.ne(t2).long().sum(), 0)

    def assertHypoEqual(self, h1, h2):
        "Check two hypos are equal"
        self.assertTensorEqual(h1["tokens"], h2["tokens"])
        self.assertAlmostEqual(h1["positional_scores"], h2["positional_scores"])
        self.assertLess(abs(h1["score"] - h2["score"]), 1e-6)
        self.assertAlmostEqual(h1["attention"], h2["attention"])

    def _test_save_and_load(self, scripted_module):
        with tempfile.NamedTemporaryFile() as f:
            scripted_module.save(f.name)
            torch.jit.load(f.name)


class TestJitSequeneceGenerator(TestJitSequenceGeneratorBase):

    @unittest.skipIf(
        torch.__version__ < "1.5.0", "Targeting OSS scriptability for the 1.5 release"
    )
    def test_export_transformer(self):
        model = self.transformer_model
        torch.jit.script(model)

    @unittest.skipIf(
        torch.__version__ < "1.6.0", "Targeting OSS scriptability for the 1.6 release"
    )
    def test_ensemble_sequence_generator(self):
        model = self.transformer_model
        generator = SequenceGenerator(
            [model], self.task.tgt_dict, beam_size=2, no_repeat_ngram_size=2
        )
        scripted_model = torch.jit.script(generator)
        self._test_save_and_load(scripted_model)


class TestJitEnsemble(TestJitSequenceGeneratorBase):

    @unittest.skipIf(
        torch.__version__ < "1.6.0", "Targeting OSS scriptability for the 1.6 release"
    )
    def test_export_ensemble_model(self):
        model = self.transformer_model
        ensemble_models = EnsembleModel([model])
        torch.jit.script(ensemble_models)


class TestExportSearch(unittest.TestCase):
    def setUp(self):
        task, _ = get_dummy_task_and_parser()
        self.tgt_dict = task.tgt_dict
        self.min_top1_prob = 0.4

    def test_export_diverse_bs(self):
{'fraction_non_alphanumeric': 0.1045593834153589, 'fraction_numerical': 0.014207447302442194, 'mean_word_length': 3.362089914945322, 'pattern_counts': {'":': 6, '<': 12, '<?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': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7638058', 'n_tokens_mistral': 3508, 'n_tokens_neox': 3373, 'n_words': 722}