Update tunes api

examples/user/prompt_tuning/classification.py

@@ -145,6 +145,13 @@ def get_creds():
             )
             time.sleep(5)
 
+            print("~~~~~~~ Downloading tuned model assets ~~~~~")
+            to_download_assets = input("Download tuned model assets? (y/N):\n")
+            if to_download_assets == "y":
+                tuned_model.download()
+
+            time.sleep(5)
+
             print("~~~~~~~ Deleting a tuned model ~~~~~")
             to_delete = input("Delete this model? (y/N):\n")
             if to_delete == "y":

examples/user/prompt_tuning/summarization.py

@@ -135,6 +135,15 @@ def get_creds():
                 tune_get_result,
             )
 
+            time.sleep(5)
+
+            print("~~~~~~~ Downloading tuned model assets ~~~~~")
+            to_download_assets = input("Download tuned model assets? (y/N):\n")
+            if to_download_assets == "y":
+                tuned_model.download()
+
+            time.sleep(5)
+
             print("~~~~~~~ Deleting a tuned model ~~~~~")
             to_delete = input("Delete this model? (y/N):\n")
             if to_delete == "y":

examples/user/prompt_tuning/tune_manager_functions.py

@@ -0,0 +1,76 @@
+import os
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+from genai.credentials import Credentials
+from genai.schemas.tunes_params import (
+    CreateTuneHyperParams,
+    CreateTuneParams,
+    DownloadAssetsParams,
+    TunesListParams,
+)
+from genai.services import FileManager, TuneManager
+
+load_dotenv()
+API_KEY = os.getenv("GENAI_KEY", None)
+ENDPOINT = os.getenv("GENAI_API", None)
+
+creds = Credentials(api_key=API_KEY, api_endpoint=ENDPOINT)
+
+
+print("\n=============== UPLOAD LOCAL FILE FOR TUNING ================")
+file_path = str(Path(__file__).parent / "file_to_tune.jsonl")
+print(" {} \n".format(file_path))
+
+file_uploaded = FileManager.upload_file(credentials=creds, file_path=file_path, purpose="tune")
+print("File uploaded: \n", file_uploaded)
+
+
+print("\n======================== CREATE TUNE ========================")
+file_ids = [file_uploaded.id]
+hyperparams = CreateTuneHyperParams(verbalizer='classify { "red", "yellow" } Input: {{input}} Output:')
+
+params = CreateTuneParams(
+    name="Tune Manager Classification",
+    model_id="google/flan-t5-xl",
+    method_id="mpt",
+    task_id="classification",
+    training_file_ids=file_ids,
+    parameters=hyperparams,
+)
+
+tune_created = TuneManager.create_tune(credentials=creds, params=params)
+print("\nNew tune: \n", tune_created)
+
+
+print("\n======================= GET TUNE BY ID ======================")
+t = tune_created.id
+tune_get = TuneManager.get_tune(credentials=creds, tune_id=t)
+print("\nGet tune result: \n", tune_get)
+
+print("\n======================== LIST TUNES =========================")
+list_params = TunesListParams(limit=5, offset=0)
+tune_list = TuneManager.list_tunes(credentials=creds, params=list_params)
+for t in tune_list.results:
+    print("\nTune ID:", t.id, ", Tune Name:", t.name)
+
+
+print("\n================  LIST EXISTING TUNE METHODS ================")
+tune_methods = TuneManager.get_tune_methods(credentials=creds)
+print("\n Available Tune Methods:\n")
+for t in tune_methods.results:
+    print("Method ID:", t.id, ", Method name:", t.name)
+
+
+print("\n===================  DOWNLOAD TUNE ASSETS ===================")
+# content can be: logs or enconder. The download will only be available when the tune is complete.
+tune = tune_list.results[0].id
+assets_params = DownloadAssetsParams(id=tune, content="encoder")
+tune_assets = TuneManager.download_tune_assets(credentials=creds, params=assets_params)
+print("\n Tune assets:", tune_assets)
+
+
+print("\n======================== DELETE TUNE ========================")
+tune_delete = TuneManager.delete_tune(credentials=creds, tune_id=tune)
+print("\nDelete tune response: \n", tune_delete)

examples/user/prompt_tuning/upload_file_create_tune.py

@@ -0,0 +1,44 @@
+import os
+from pathlib import Path
+
+from dotenv import load_dotenv
+
+from genai.credentials import Credentials
+from genai.schemas.tunes_params import CreateTuneParams
+from genai.services import FileManager, TuneManager
+
+load_dotenv()
+API_KEY = os.getenv("GENAI_KEY", None)
+ENDPOINT = os.getenv("GENAI_API", None)
+
+creds = Credentials(api_key=API_KEY, api_endpoint=ENDPOINT)
+
+# # UPLOAD FILE
+# file_path = "<path-to-file>"
+print("======================== UPLOAD LOCAL FILE ========================")
+file_path = str(Path(__file__).parent / "file_to_tune.jsonl")
+print(" {} \n".format(file_path))
+
+upload_file = FileManager.upload_file(credentials=creds, file_path=file_path, purpose="tune")
+print("File uploaded: \n", upload_file)
+
+print("\n===================== GET UPLOADED FILE ID ========================")
+file_list = FileManager.list_files(credentials=creds)
+for f in file_list.results:
+    if f.file_name == "file_to_tune.jsonl":
+        file_id = f.id
+        break
+print("Uploaded file has id = {}\n".format(file_id))
+
+print("\n================= CREATE TUNE FOR GENERATION TASK ==================")
+
+tunes_params = CreateTuneParams(
+    name="flan-t5-xl",
+    model_id="google/flan-t5-xl",
+    method_id="pt",
+    task_id="generation",
+    training_file_ids=[file_id],
+)
+
+tune_create = TuneManager.create_tune(credentials=creds, params=tunes_params)
+print("Created tune has the id: \n ", tune_create.id)

src/genai/model.py

@@ -24,6 +24,7 @@
 from genai.schemas.tunes_params import (
     CreateTuneHyperParams,
     CreateTuneParams,
+    DownloadAssetsParams,
     TunesListParams,
 )
 from genai.services import AsyncResponseGenerator, ServiceInterface
@@ -367,6 +368,12 @@ def delete(self):
             raise GenAiException(ValueError("Tuned model not found. Currently method supports only tuned models."))
         TuneManager.delete_tune(service=self.service, tune_id=self.model)
 
+    def download(self):
+        enconder_params = DownloadAssetsParams(id=self.model, content="encoder")
+        logs_params = DownloadAssetsParams(id=self.model, content="logs")
+        TuneManager.download_tune_assets(service=self.service, params=enconder_params)
+        TuneManager.download_tune_assets(service=self.service, params=logs_params)
+
     @staticmethod
     def models(credentials: Credentials = None, service: ServiceInterface = None) -> list[ModelCard]:
         """Get a list of models

src/genai/routers/tunes.py

@@ -1,5 +1,9 @@
 from genai.exceptions import GenAiException
-from genai.schemas.tunes_params import CreateTuneParams, TunesListParams
+from genai.schemas.tunes_params import (
+    CreateTuneParams,
+    DownloadAssetsParams,
+    TunesListParams,
+)
 from genai.services.request_handler import RequestHandler
 from genai.utils.request_utils import sanitize_params
 
@@ -73,3 +77,27 @@ def delete_tune(self, tune_id: str):
             return RequestHandler.delete(endpoint, key=self.key, parameters=tune_id)
         except Exception as e:
             raise GenAiException(e)
+
+    def get_tune_methods(self):
+        """Get list of tune methods.
+
+        Returns:
+            Any: json with info about the available tune methods.
+        """
+        try:
+            endpoint = self.service_url + "/tune_methods"
+            return RequestHandler.get(endpoint, key=self.key)
+        except Exception as e:
+            raise GenAiException(e)
+
+    def download_tune_assets(self, params: DownloadAssetsParams):
+        """Download tune asset.
+
+        Returns:
+            Any: json with info about the downloaded tune asset.
+        """
+        try:
+            endpoint = self.service_url + TunesRouter.TUNES + "/" + params.id + "/content/" + params.content
+            return RequestHandler.get(endpoint, key=self.key)
+        except Exception as e:
+            raise GenAiException(e)

src/genai/schemas/descriptions.py

@@ -69,6 +69,8 @@ class TunesAPIDescriptions:
     STATUS = "Filters the items to be returned based on their status. Possible values are: INITIALIZING, NOT_STARTED, PENDING, HALTED, RUNNING, QUEUED, COMPLETED, FAILED."
     INIT_METHOD = "Initialization method to be used. Possible values are RANDOM or TEXT. Defaults to RANDOM. Used only if the method_id is 'pt' = Prompt Tuning."
     INIT_TEXT = "Initialization text to be used. This is only applicable if init_method == TEXT. Used only if the method_id is 'pt' = Prompt Tuning."
+    ID = "The ID of the tune."
+    CONTENT = "The name of the asset. Available options are encoder and logs."
 
 
 class FilesAPIDescriptions:

src/genai/schemas/responses.py

@@ -229,6 +229,62 @@ class TuneGetResponse(GenAiResponseModel):
     results: Optional[TuneInfoResult]
 
 
+class TuneMethodsInfo(GenAiResponseModel):
+    id: str
+    name: str
+
+
+class TuneMethodsGetResponse(GenAiResponseModel):
+    results: Optional[List[TuneMethodsInfo]]
+
+
+class FileFormatResult(GenAiResponseModel):
+    id: int
+    name: str
+
+
+class FileInfoResult(GenAiResponseModel):
+    id: str
+    bytes: str
+    file_name: str
+    purpose: str
+    storage_provider_location: Optional[str]
+    created_at: datetime
+    file_formats: List[FileFormatResult]
+
+
+class TuneParameters(GenAiResponseModel):
+    accumulate_steps: Optional[int]
+    batch_size: Optional[int]
+    learning_rate: Optional[float]
+    max_input_tokens: Optional[int]
+    max_output_tokens: Optional[int]
+    num_epochs: Optional[int]
+    num_virtual_tokens: Optional[int]
+    verbalizer: Optional[str]
+
+
+class TuneInfoResult(GenAiResponseModel):
+    id: str
+    name: str
+    model_id: str
+    model_name: str
+    method_id: Optional[str]
+    method_name: Optional[str]
+    status: str
+    task_id: str
+    task_name: Optional[str]
+    parameters: Optional[TuneParameters]
+    created_at: datetime
+    preferred: Optional[bool]
+    datapoints: Optional[dict]
+    validation_files: Optional[list]
+    training_files: Optional[list]
+    evaluation_files: Optional[list]
+    status_message: Optional[str]
+    started_at: Optional[datetime]
+
+
 class ModelCard(GenAiResponseModel):
     id: Optional[str]
     name: Optional[str]

src/genai/schemas/tunes_params.py

@@ -52,3 +52,14 @@ class Config:
     training_file_ids: list[str] = Field(None, description=tx.TRAINING_FILE_IDS)
     validation_file_ids: Optional[list[str]] = Field(None, description=tx.VALIDATION_FILE_IDS)
     parameters: Optional[CreateTuneHyperParams] = Field(None, description=tx.PARAMETERS)
+
+
+class DownloadAssetsParams(BaseModel):
+    """Class to hold the parameters for downloading tune assets."""
+
+    class Config:
+        anystr_strip_whitespace: True
+        # extra: Extra.forbid
+
+    id: str = Field(None, description=tx.ID)
+    content: str = Field(None, description=tx.CONTENT)