firebase · kevinthecheung · Sep 15, 2020 · Aug 24, 2020 · Aug 26, 2020 · Sep 15, 2020
diff --git a/machine-learning/README.md b/machine-learning/README.md
@@ -57,18 +57,31 @@ Also see the [Firebase ML API Tutorial Colab/Jupyter notebook][colab].
 
 ```
 $ ./manage-ml.py list
-fish_detector    8716935   vision
-barcode_scanner  8716959   vision
-smart_reply      8716981   natural_language
-$ ./manage-ml.py new ~/yak.tflite yak_detector --tags vision,experimental
-Uploading model to Cloud Storage...
+         Name              ID            Tags
+---------------------- ---------- ------------------
+ fish_recognizer        12990533   vision
+ barcode_scanner        12990544   vision
+$ ./manage-ml.py new yak_detector -f model.tflite -t vision,experimental
+Uploading to Cloud Storage...
 Model uploaded and published:
-yak_detector     8717019   experimental, vision
-$ ./manage-ml.py update 8717019 --remove_tags experimental
-$ ./manage-ml.py delete 8716959
+ yak_detector   12990577   experimental, vision
+$ ./manage-ml.py new flower_classifier -a projects/12345/locations/us-central1/models/ICN12345
+Model uploaded and published:
+ flower_classifier   12990597
+$ ./manage-ml.py list
+         Name              ID              Tags
+---------------------- ---------- ----------------------
+ fish_recognizer        12990533   vision
+ barcode_scanner        12990544   vision
+ yak_detector           12990577   experimental, vision
+ flower_classifier      12990597
+$ ./manage-ml.py update 12990577 --remove_tags experimental
+$ ./manage-ml.py delete 12990544
 $ ./manage-ml.py list
-fish_detector    8716935   vision
-smart_reply      8716981   natural_language
-yak_detector     8717019   vision
+         Name              ID            Tags
+---------------------- ---------- ------------------
+ fish_recognizer        12990533   vision
+ yak_detector           12990577   vision
+ flower_classifier      12990597
 $
 ```
diff --git a/machine-learning/manage-ml.py b/machine-learning/manage-ml.py
@@ -42,6 +42,25 @@ def upload_model(model_file, name, tags=None):
   print_models([new_model], headers=False)
 
 
+def add_automl_model(model_ref, name, tags=None):
+  """Add an AutoML tflite model file to the project and publish it."""
+  # Create the model object
+  model_source = ml.TFLiteAutoMlSource(model_ref)
+  model = ml.Model(
+      display_name=name,
+      model_format=ml.TFLiteFormat(model_source=model_source))
+  if tags is not None:
+    model.tags = tags
+
+  # Add the model to your Firebase project and publish it
+  new_model = ml.create_model(model)
+  new_model.wait_for_unlocked()
+  ml.publish_model(new_model.model_id)
+
+  print('Model uploaded and published:')
+  print_models([new_model], headers=False)
+
+
 def list_models(filter_exp=''):
   """List the models in the project."""
   models = ml.list_models(list_filter=filter_exp).iterate_all()
@@ -123,11 +142,17 @@ def main():
       dest='command', required=True, metavar='command')
 
   new_parser = subparsers.add_parser(
-      'new', help='upload a tflite model to your project')
-  new_parser.add_argument(
-      'model_file', type=str, help='path to the tflite file')
+      'new', help='upload a tflite model file or AutoML model reference to'
+                  ' your project')
   new_parser.add_argument(
       'name', type=str, help='display name for the new model')
+  new_source_group = new_parser.add_mutually_exclusive_group(required=True)
+  new_source_group.add_argument(
+      '-f', '--file', type=str, help='path to the tflite file')
+  new_source_group.add_argument(
+      '-a', '--automl', type=str, help='AutoML model reference (e.g. projects/'
+                                       '12345678/locations/us-central1/models/'
+                                       'ICN1234567890)')
   new_parser.add_argument(
       '-t', '--tags', type=str, help='comma-separated list of tags')
 
@@ -165,9 +190,12 @@ def main():
 
   args = main_parser.parse_args()
   try:
-    if args.command == 'new':
+    if args.command == 'new' and args.file is not None:
+      tags = args.tags.split(',') if args.tags is not None else None
+      upload_model(args.file.strip(), args.name.strip(), tags)
+    if args.command == 'new' and args.automl is not None:
       tags = args.tags.split(',') if args.tags is not None else None
-      upload_model(args.model_file, args.name, tags)
+      add_automl_model(args.automl.strip(), args.name.strip(), tags)
     elif args.command == 'list':
       list_models(args.filter)
     elif args.command == 'info':

diff --git a/machine-learning/requirements.txt b/machine-learning/requirements.txt
@@ -1,2 +1,2 @@
 beautifultable~=1.0.0
-firebase-admin~=4.3.0
+firebase-admin~=4.4.0