def model_list(): ml_model_1 = ml.Model(display_name=_random_identifier('TestModel123_list1_')) model_1 = ml.create_model(model=ml_model_1) ml_model_2 = ml.Model(display_name=_random_identifier('TestModel123_list2_'), tags=['test_tag123']) model_2 = ml.create_model(model=ml_model_2) yield [model_1, model_2] _clean_up_model(model_1) _clean_up_model(model_2)
def automl_model(): assert _AUTOML_ENABLED # It takes > 20 minutes to train a model, so we expect a predefined AutoMl # model named 'admin_sdk_integ_test1' to exist in the project, or we skip # the test. automl_client = automl_v1.AutoMlClient() project_id = firebase_admin.get_app().project_id parent = automl_client.location_path(project_id, 'us-central1') models = automl_client.list_models(parent, filter_="display_name=admin_sdk_integ_test1") # Expecting exactly one. (Ok to use last one if somehow more than 1) automl_ref = None for model in models: automl_ref = model.name # Skip if no pre-defined model. (It takes min > 20 minutes to train a model) if automl_ref is None: pytest.skip("No pre-existing AutoML model found. Skipping test") source = ml.TFLiteAutoMlSource(automl_ref) tflite_format = ml.TFLiteFormat(model_source=source) ml_model = ml.Model( display_name=_random_identifier('TestModel_automl_'), tags=['test_automl'], model_format=tflite_format) model = ml.create_model(model=ml_model) yield model _clean_up_model(model)
def test_model_as_dict_for_upload(self): model_source = ml.TFLiteGCSModelSource(gcs_tflite_uri=GCS_TFLITE_URI) model_format = ml.TFLiteFormat(model_source=model_source) model = ml.Model(display_name=DISPLAY_NAME_1, model_format=model_format) assert model.as_dict(for_upload=True) == { 'displayName': DISPLAY_NAME_1, 'tfliteModel': { 'gcsTfliteUri': GCS_TFLITE_SIGNED_URI } }
def test_model_format_source_creation(self): model_source = ml.TFLiteGCSModelSource(gcs_tflite_uri=GCS_TFLITE_URI) model_format = ml.TFLiteFormat(model_source=model_source) model = ml.Model(display_name=DISPLAY_NAME_1, model_format=model_format) assert model.as_dict() == { 'displayName': DISPLAY_NAME_1, 'tfliteModel': { 'gcsTfliteUri': GCS_TFLITE_URI } }
def firebase_model(request): args = request.param tflite_format = None file_name = args.get('file_name') if file_name: file_path = testutils.resource_filename(file_name) source = ml.TFLiteGCSModelSource.from_tflite_model_file(file_path) tflite_format = ml.TFLiteFormat(model_source=source) ml_model = ml.Model(display_name=args.get('display_name'), tags=args.get('tags'), model_format=tflite_format) model = ml.create_model(model=ml_model) yield model _clean_up_model(model)
def test_from_keras_model(keras_model): source = ml.TFLiteGCSModelSource.from_keras_model(keras_model, 'model2.tflite') assert re.search( '^gs://.*/Firebase/ML/Models/model2.tflite$', source.gcs_tflite_uri) is not None # Validate the conversion by creating a model model_format = ml.TFLiteFormat(model_source=source) model = ml.Model(display_name=_random_identifier('KerasModel_'), model_format=model_format) created_model = ml.create_model(model) try: check_model(created_model, {'display_name': model.display_name}) check_tflite_gcs_format(created_model) finally: _clean_up_model(created_model)
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 test_from_saved_model(saved_model_dir): # Test the conversion helper source = ml.TFLiteGCSModelSource.from_saved_model(saved_model_dir, 'model3.tflite') assert re.search( '^gs://.*/Firebase/ML/Models/model3.tflite$', source.gcs_tflite_uri) is not None # Validate the conversion by creating a model model_format = ml.TFLiteFormat(model_source=source) model = ml.Model(display_name=_random_identifier('SavedModel_'), model_format=model_format) created_model = ml.create_model(model) try: assert created_model.model_id is not None assert created_model.validation_error is None finally: _clean_up_model(created_model)
def upload_model(model_file, name, tags=None): """Upload a tflite model file to the project and publish it.""" # Load a tflite file and upload it to Cloud Storage print('Uploading to Cloud Storage...') model_source = ml.TFLiteGCSModelSource.from_tflite_model_file(model_file) # Create the model object tflite_format = ml.TFLiteFormat(model_source=model_source) model = ml.Model(display_name=name, model_format=tflite_format) if tags is not None: model.tags = tags # Add the model to your Firebase project and publish it new_model = ml.create_model(model) ml.publish_model(new_model.model_id) print('Model uploaded and published:') print_models([new_model], headers=False)
def test_model_keyword_based_creation_and_setters(self): model = ml.Model(display_name=DISPLAY_NAME_1, tags=TAGS, model_format=TFLITE_FORMAT) assert model.display_name == DISPLAY_NAME_1 assert model.tags == TAGS assert model.model_format == TFLITE_FORMAT assert model.as_dict() == { 'displayName': DISPLAY_NAME_1, 'tags': TAGS, 'tfliteModel': TFLITE_FORMAT_JSON } model.display_name = DISPLAY_NAME_2 model.tags = TAGS_2 model.model_format = TFLITE_FORMAT_2 assert model.as_dict() == { 'displayName': DISPLAY_NAME_2, 'tags': TAGS_2, 'tfliteModel': TFLITE_FORMAT_JSON_2 }
def test_wait_for_unlocked_not_locked(self): model = ml.Model(display_name="not_locked") model.wait_for_unlocked()
def test_model_format_validation_errors(self, model_format): with pytest.raises(TypeError) as excinfo: ml.Model(model_format=model_format) check_error(excinfo, TypeError, 'Model format must be a ModelFormat object.')
def test_model_tags_validation_errors(self, tags, exc_type, error_message): with pytest.raises(exc_type) as excinfo: ml.Model(tags=tags) check_error(excinfo, exc_type, error_message)
def test_model_display_name_validation_errors(self, display_name, exc_type): with pytest.raises(exc_type) as excinfo: ml.Model(display_name=display_name) check_error(excinfo, exc_type)