Iterate on torch.hub interface

Summary: Pull Request resolved: https://github.com/pytorch/fairseq/pull/793

Differential Revision: D15758755

Pulled By: myleott

fbshipit-source-id: b93e4ac11bde36a0b59b4d6d1c84d31c3124d767

hubconf.py

@@ -5,48 +5,11 @@
 # 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 fairseq.models.transformer import TransformerModel
-from fairseq.models.fconv import FConvModel
-from fairseq.models.fconv_self_att import FConvModelSelfAtt
-from generator import Generator
-from fairseq import options
+from fairseq.models import MODEL_REGISTRY
 
 
-dependencies = [
-    'torch',
-    'sacremoses',
-    'subword_nmt',
-]
+dependencies = ['torch']
 
 
-def transformer(*args, **kwargs):
-    """
-    Transformer model from `"Attention Is All You Need" (Vaswani, et al, 2017)
-    <https://arxiv.org/abs/1706.03762>`_.
-    """
-    parser = options.get_interactive_generation_parser()
-    model = TransformerModel.from_pretrained(parser, *args, **kwargs)
-    return model
-
-
-def fconv(*args, **kwargs):
-    """
-    A fully convolutional model, i.e. a convolutional encoder and a
-    convolutional decoder, as described in `"Convolutional Sequence to Sequence
-    Learning" (Gehring et al., 2017) <https://arxiv.org/abs/1705.03122>`_.
-    """
-    parser = options.get_interactive_generation_parser()
-    model = FConvModel.from_pretrained(parser, *args, **kwargs)
-    return model
-
-
-def fconv_self_att(*args, **kwargs):
-    parser = options.get_interactive_generation_parser()
-    model = FConvModelSelfAtt.from_pretrained(parser, *args, **kwargs)
-    return model
-
-
-def generator(*args, **kwargs):
-    parser = options.get_generation_parser(interactive=True)
-    generator = Generator.from_pretrained(parser, *args, **kwargs)
-    return generator
+for model, cls in MODEL_REGISTRY.items():
+    globals()[model] = cls.from_pretrained