Spaces:
Running
Running
File size: 2,896 Bytes
650c5f6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 |
# 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 unittest
from argparse import ArgumentParser
from dataclasses import dataclass, field
from fairseq.dataclass import FairseqDataclass
from fairseq.dataclass.utils import gen_parser_from_dataclass
@dataclass
class A(FairseqDataclass):
data: str = field(default="test", metadata={"help": "the data input"})
num_layers: int = field(default=200, metadata={"help": "more layers is better?"})
@dataclass
class B(FairseqDataclass):
bar: A = field(default=A())
foo: int = field(default=0, metadata={"help": "not a bar"})
@dataclass
class D(FairseqDataclass):
arch: A = field(default=A())
foo: int = field(default=0, metadata={"help": "not a bar"})
@dataclass
class C(FairseqDataclass):
data: str = field(default="test", metadata={"help": "root level data input"})
encoder: D = field(default=D())
decoder: A = field(default=A())
lr: int = field(default=0, metadata={"help": "learning rate"})
class TestDataclassUtils(unittest.TestCase):
def test_argparse_convert_basic(self):
parser = ArgumentParser()
gen_parser_from_dataclass(parser, A(), True)
args = parser.parse_args(["--num-layers", '10', "the/data/path"])
self.assertEqual(args.num_layers, 10)
self.assertEqual(args.data, "the/data/path")
def test_argparse_recursive(self):
parser = ArgumentParser()
gen_parser_from_dataclass(parser, B(), True)
args = parser.parse_args(["--num-layers", "10", "--foo", "10", "the/data/path"])
self.assertEqual(args.num_layers, 10)
self.assertEqual(args.foo, 10)
self.assertEqual(args.data, "the/data/path")
def test_argparse_recursive_prefixing(self):
self.maxDiff = None
parser = ArgumentParser()
gen_parser_from_dataclass(parser, C(), True, "")
args = parser.parse_args(
[
"--encoder-arch-data",
"ENCODER_ARCH_DATA",
"--encoder-arch-num-layers",
"10",
"--encoder-foo",
"10",
"--decoder-data",
"DECODER_DATA",
"--decoder-num-layers",
"10",
"--lr",
"10",
"the/data/path",
]
)
self.assertEqual(args.encoder_arch_data, "ENCODER_ARCH_DATA")
self.assertEqual(args.encoder_arch_num_layers, 10)
self.assertEqual(args.encoder_foo, 10)
self.assertEqual(args.decoder_data, "DECODER_DATA")
self.assertEqual(args.decoder_num_layers, 10)
self.assertEqual(args.lr, 10)
self.assertEqual(args.data, "the/data/path")
if __name__ == "__main__":
unittest.main()
|