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)
