Add save-pytext-snapshot command to PyText cmdline

pytext/main.py

@@ -38,6 +38,7 @@
     export_saved_model_to_torchscript,
     get_logits as workflow_get_logits,
     prepare_task_metadata,
+    save_pytext_snapshot as workflow_save_pytext_snapshot,
     test_model_from_snapshot_path,
     train_model,
 )
@@ -480,5 +481,16 @@ def get_logits(context, model_snapshot, test_path, use_cuda, output_path, field_
     workflow_get_logits(model_snapshot, use_cuda, output_path, test_path, field_names)
 
 
+@main.command()
+@click.pass_context
+def save_pytext_snapshot(context):
+    """Load a PyText task and save snapshot for later use.
+    This is helpful when you want to plug in a pretrained encoder in a PyText
+    task and either test or generate logits using the task.
+    """
+    config = context.obj.load_config()
+    workflow_save_pytext_snapshot(config)
+
+
 if __name__ == "__main__":
     main()

pytext/workflow.py

@@ -348,6 +348,20 @@ def get_logits(
                     )
 
 
+def save_pytext_snapshot(config: PyTextConfig,) -> None:
+    task, training_state = prepare_task(
+        config,
+        dist_init_url=None,
+        device_id=0,
+        rank=0,
+        world_size=1,
+        metric_channels=None,
+        metadata=None,
+    )
+    print("Task set up successful.\n")
+    save_and_export(config, task)
+
+
 def dict_zip(*dicts, value_only=False):
     dict_keys = dicts[0].keys()
     return (