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 toNone
.subcategory (
str
, optional): The subcategory to filter models by. If provided,category
must also be provided. Defaults toNone
.
Returns:
List[Model]
A list ofModel
instances matching the criteria.Returns an empty list if no models are found.
Raises:
ValueError
: If theprovider
isNone
or invalid.ValueError
: Ifsubcategory
is provided withoutcategory
.
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, otherwiseNone
.
Raises:
ValueError
: Ifprovider
ormodel_name
isNone
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 typeHuggingFace Model
andCustom Model
.For
Custom Models
that are not self-hosted, either the path to the model file hasupload_file_path
has to be provided or a downloadable URL containing the modelupload_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 ofModelConfig
containing all the necessary details for the model to be saved. This object should be created using one of the specific builders likeModelConfig.huggingface_builder()
orModelConfig.custom_builder()
.
Returns:
Model
: An instance ofModel
class representing the successfully saved model. This object can be used for further operations.
Raises:
ValueError
: Ifmodel_config
isNone
, not an instance ofModelConfig
, or if essential provider-specific fields are missing (e.g.,hf_model_name
for Hugging Face models, orhosted_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 ofModelConfig
representing the model to be requested. It should be created using theModelConfig.vertexai_builder()
for Vertex AI models. The config must contain arequest_model
attribute.
Returns:
Model
: An instance of theModel
class representing the requested model. This object can be used for further operations.
Raises:
ValueError
: Ifmodel_config
isNone
or not an instance ofModelConfig
.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)