diff --git a/clients/python/llmengine/data_types.py b/clients/python/llmengine/data_types.py index baa9e087..e17e755a 100644 --- a/clients/python/llmengine/data_types.py +++ b/clients/python/llmengine/data_types.py @@ -383,6 +383,16 @@ class CreateFineTuneRequest(BaseModel): ) """Hyperparameters to pass in to training job.""" + wandb_config: Optional[Dict[str, Any]] = Field( + default=None, description="Configuration for Weights and Biases." + ) + """ + A dict of configuration parameters for Weights & Biases. See [Weights & Biases](https://docs.wandb.ai/ref/python/init) for more information. + Set `hyperparameter["report_to"]` to `wandb` to enable automatic finetune metrics logging. + Must include `api_key` field which is the wandb API key. + Also supports setting `base_url` to use a custom Weights & Biases server. + """ + suffix: Optional[str] = Field( default=None, description="Optional user-provided identifier suffix for the fine-tuned model.", diff --git a/clients/python/llmengine/fine_tuning.py b/clients/python/llmengine/fine_tuning.py index c280ac39..d0a0ef86 100644 --- a/clients/python/llmengine/fine_tuning.py +++ b/clients/python/llmengine/fine_tuning.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional, Union +from typing import Any, Dict, Optional, Union from llmengine.api_engine import DEFAULT_TIMEOUT, APIEngine from llmengine.data_types import ( @@ -29,6 +29,7 @@ def create( training_file: str, validation_file: Optional[str] = None, hyperparameters: Optional[Dict[str, Union[str, int, float]]] = None, + wandb_config: Optional[Dict[str, Any]] = None, suffix: Optional[str] = None, ) -> CreateFineTuneResponse: """ @@ -66,6 +67,12 @@ def create( * `epochs`: Number of fine-tuning epochs. This should be less than 20. (Default: 5) * `weight_decay`: Regularization penalty applied to learned weights. (Default: 0.001) + wandb_config (`Optional[Dict[str, Any]]`): + A dict of configuration parameters for Weights & Biases. See [Weights & Biases](https://docs.wandb.ai/ref/python/init) for more information. + Set `hyperparameter["report_to"]` to `wandb` to enable automatic finetune metrics logging. + Must include `api_key` field which is the wandb API key. + Also supports setting `base_url` to use a custom Weights & Biases server. + suffix (`Optional[str]`): A string that will be added to your fine-tuned model name. If present, the entire fine-tuned model name will be formatted like `"[model].[suffix].[YYYY-MM-DD-HH-MM-SS]"`. If absent, the @@ -134,6 +141,7 @@ def create( training_file=training_file, validation_file=validation_file, hyperparameters=hyperparameters, + wandb_config=wandb_config, suffix=suffix, ) response = cls.post_sync( diff --git a/docs/model_zoo.md b/docs/model_zoo.md index da07c287..264196a6 100644 --- a/docs/model_zoo.md +++ b/docs/model_zoo.md @@ -10,6 +10,7 @@ Scale hosts the following models in the LLM Engine Model Zoo: | `llama-2-13b` | ✅ | | | `llama-2-13b-chat` | ✅ | | | `llama-2-70b` | ✅ | | +| `llama-2-70b-chat` | ✅ | | | `falcon-7b` | ✅ | | | `falcon-7b-instruct` | ✅ | | | `falcon-40b` | ✅ | |