Model Repo

The ModelRepo class acts as a repository for managing machine learning models within BerryDB. It provides methods to:

  • Retrieve existing models from the repository (get).

  • Save new model configurations (save).

  • Request models, particularly for providers like Vertex AI (request).

ModelRepo.vertexai_models_for_request()

Fetches a list of models the user can request from the Vertex AI API.

ModelRepo.get(provider, category=None, subcategory=None)

Fetches a list of models from the model repository based on the provider and optional category/subcategory filters.

Parameters:

  • provider (ModelProvider): The provider of the model(s) to fetch (e.g., ModelProvider.HUGGING_FACE_MODEL, ModelProvider.VERTEX_AI_MODEL). This parameter is required.

  • category (str, optional): The category to filter models by. Defaults to None.

  • subcategory (str, optional): The subcategory to filter models by. If provided, category must also be provided. Defaults to None.

Returns:

  • List[Model]A list of Model instances matching the criteria.

    Returns an empty list if no models are found.

Raises:

  • ValueError: If the provider is None or invalid.

  • ValueError: If subcategory is provided without category.

Example:

from berrydb import BerryDB, ModelProvider

# Initialize BerryDB (if not already done)
# BerryDB.init("your-berrydb-host")

berrydb_api_key = "BERRYDB_API_KEY"

# Get the model repository
repo = BerryDB.model_repo(berrydb_api_key)

# Example 1: Get all Custom models
try:
    custom_models_list = repo.get(provider=ModelProvider.CUSTOM_MODEL)
    if custom_models_list:
        print(f"Fetched {len(custom_models_list)} Custom models:")
        for model_item in custom_models_list:
            print(f"- {model_item.config.name}")
    else:
        print("No Custom models found.")
except Exception as e:
    print(f"Error fetching Custom models: {e}")

# Example 2: Get Hugging Face models in 'Text Generation' category
try:
    text_gen_models = repo.get(
        provider=ModelProvider.HUGGING_FACE_MODEL,
        category="Text Generation" # Make sure this category exists
    )
    if text_gen_models:
        print(f"Found {len(text_gen_models)} Text Generation models.")
    else:
        print("No Text Generation models found for Hugging Face.")
except Exception as e:
    print(f"Error fetching Text Generation models: {e}")

# Example 3: Get Custom models in 'Computer Vision' category and 'Image Classification' subcategory
try:
    image_classification_models = repo.get(
        provider=ModelProvider.CUSTOM_MODEL,
        category="Computer Vision", # Make sure this category exists
        subcategory="Image Classification" # Make sure this subcategory exists for the category
    )
    if image_classification_models:
        print(f"Found {len(image_classification_models)} Image Classification models.")
    else:
        print("No Image Classification models found for Custom provider with specified category/subcategory.")
except Exception as e:
    print(f"Error fetching Image Classification models: {e}")
ModelRepo.get_by_name(provider, model_name)

Fetches a specific model by its name and provider.

Parameters:

  • provider (ModelProvider): The provider of the model to fetch.

  • model_name (str): The name of the specific model to fetch.

Returns:

  • Model | None: An instance of the Model class if found, otherwise None.

Raises:

  • ValueError: If provider or model_name is None or invalid.

  • An exception (via Utils.handleApiCallFailure) if another API error occurs (e.g. server error).

Example:

from berrydb import BerryDB, ModelProvider

# Initialize BerryDB (if not already done)
# BerryDB.init("your-berrydb-host")

berrydb_api_key = "BERRYDB_API_KEY"
repo = BerryDB.model_repo(berrydb_api_key)

try:
    # Example 1: Get a specific Hugging Face model
    hf_model = repo.get_by_name(
        provider=ModelProvider.HUGGING_FACE_MODEL,
        model_name="my-sentiment-analyzer"
    )
    if hf_model:
        print(f"Fetched Hugging Face model: {hf_model.config.name}")
    else:
        print("Hugging Face model 'my-sentiment-analyzer' not found.")

    # Example 2: Get a specific Vertex AI model
    vertex_model = repo.get_by_name(
        provider=ModelProvider.VERTEX_AI_MODEL,
        model_name="gemini-pro" # Ensure this model exists
    )
    if vertex_model:
        print(f"Fetched Vertex AI model: {vertex_model.config.name}")
    else:
        print("Vertex AI model 'gemini-pro' not found.")

except Exception as e: # Catch other potential API errors
    print(f"An error occurred: {e}")
ModelRepo.save(model_config)

Saves a new model to BerryDB.

This method takes a ModelConfig object and registers the new model with BerryDB. You may save models of type HuggingFace Model and Custom Model.

For Custom Models that are not self-hosted, either the path to the model file has upload_file_path has to be provided or a downloadable URL containing the model upload_file_url has to be provided.

Note

If your model has a specific folder and file structure, please zip the folder accordingly and use it to add the model.

Parameters:

  • model_config (ModelConfig): An instance of ModelConfig containing all the necessary details for the model to be saved. This object should be created using one of the specific builders like ModelConfig.huggingface_builder() or ModelConfig.custom_builder().

Returns:

  • Model: An instance of Model class representing the successfully saved model. This object can be used for further operations.

Raises:

  • ValueError: If model_config is None, not an instance of ModelConfig, or if essential provider-specific fields are missing (e.g., hf_model_name for Hugging Face models, or hosted_url for self-hosted custom models, or file paths/URLs for non-self-hosted custom models).

Example:

from berrydb import BerryDB, ModelProvider, ModelConfig

berrydb_api_key = "BERRYDB_API_KEY"
repo = BerryDB.model_repo(berrydb_api_key)

# Example 1: Save a Hugging Face model
hf_config = (
    ModelConfig.huggingface_builder()
    .name("my-sentiment-classifier")
    .description("A sentiment analysis model from Hugging Face.")
    .hf_model_name("distilbert-base-uncased-finetuned-sst-2-english")
    .build(api_key="BERRYDB_API_KEY")
)
saved_hf_model = repo.save(hf_config)

# Example 2: Save a self-hosted Custom Model
custom_self_hosted_config = (
    ModelConfig.custom_builder()
    .name("my-custom-ner-api")
    .description("A self-hosted NER model.")
    .self_hosted(True)
    .hosted_url("http://my-ner-service.example.com/predict")
    .build(api_key="BERRYDB_API_KEY")
)
saved_custom_model = repo.save(custom_self_hosted_config)
ModelRepo.request(model_config)

Requests a model deployment from the specified provider.

This method initiates a request to deploy a model. Currently, it only supports Vertex AI models.

Parameters:

  • model_config (ModelConfig): An instance of ModelConfig representing the model to be requested. It should be created using the ModelConfig.vertexai_builder() for Vertex AI models. The config must contain a request_model attribute.

Returns:

  • Model: An instance of the Model class representing the requested model. This object can be used for further operations.

Raises:

  • ValueError: If model_config is None or not an instance of ModelConfig.

  • NotImplementedError: If the method is called for a provider other than Vertex AI.

Example:

from berrydb import BerryDB, ModelProvider, ModelConfig

# Initialize BerryDB (if not already done)
# BerryDB.init("your-berrydb-host")

berrydb_api_key = "BERRYDB_API_KEY"
repo = BerryDB.model_repo(berrydb_api_key)

# Create a Vertex AI model config
vertex_config = (
    ModelConfig.vertexai_builder()
    .name("gemini-pro")
    .request_model("gemini-pro")
    .notes("Requesting a Gemini Pro model.")
    .build(api_key="BERRYDB_API_KEY")
)

# Request the Vertex AI model
requested_model = repo.request(vertex_config)