diff --git a/pytorch_pretrained_bert/__main__.py b/pytorch_pretrained_bert/__main__.py index 73f1909b..79ad8429 100644 --- a/pytorch_pretrained_bert/__main__.py +++ b/pytorch_pretrained_bert/__main__.py @@ -1,5 +1,5 @@ # coding: utf8 -if __name__ == '__main__': +def main(): import sys try: from .convert_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch @@ -17,3 +17,6 @@ if __name__ == '__main__': TF_CONFIG = sys.argv.pop() TF_CHECKPOINT = sys.argv.pop() convert_tf_checkpoint_to_pytorch(TF_CHECKPOINT, TF_CONFIG, PYTORCH_DUMP_OUTPUT) + +if __name__ == '__main__': + main() diff --git a/setup.py b/setup.py index dbfeb2c6..e9b5c077 100644 --- a/setup.py +++ b/setup.py @@ -53,7 +53,11 @@ setup( 'boto3', 'requests', 'tqdm'], - scripts=["bin/pytorch_pretrained_bert"], + entry_points={ + 'console_scripts': [ + "pytorch_pretrained_bert=pytorch_pretrained_bert.__main__:main" + ] + }, python_requires='>=3.5.0', tests_require=['pytest'], classifiers=[