Update to 3.1.3 (#76)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

README.md

@@ -21,10 +21,10 @@ production:
     key: "<YOUR-KEY-HERE>"
 ```
 
-# Documentation
+## Documentation
 
 Package documentation here: [https://lamini-ai.github.io/](https://lamini-ai.github.io/)
 
-# Github
+## Github
 
 Package source code here: [https://github.com/lamini-ai/lamini](https://github.com/lamini-ai/lamini)

lamini/__init__.py

@@ -1,18 +1,5 @@
 # Turn of isort, because alphabetic order for the following imports causes circular dependency issues
 
-# isort: off
-from lamini.error import error
-
-from lamini.api.lamini import Lamini
-from lamini.api.classifier import Classifier
-from lamini.api.embedding import Embedding
-from lamini.classify.lamini_classifier import LaminiClassifier
-from lamini.generation.generation_node import GenerationNode
-from lamini.generation.generation_pipeline import GenerationPipeline
-from lamini.generation.base_prompt_object import PromptObject
-from lamini.generation.split_response_node import SplitResponseNode
-from lamini.api.streaming_completion import StreamingCompletion
-
 import os
 
 api_key = os.environ.get("LAMINI_API_KEY", None)
@@ -29,5 +16,23 @@
 batch_size = int(os.environ.get("LAMINI_BATCH_SIZE", 5))
 static_batching = bool(os.environ.get("LAMINI_STATIC_BATCHING", False))
 bypass_reservation = bool(os.environ.get("LAMINI_BYPASS_RESERVATION", False))
+gate_pipeline_batch_completions = bool(
+    os.environ.get("GATE_PIPELINE_BATCH_COMPLETIONS", False)
+)
 
-__version__ = "3.1.0"
+__version__ = "3.1.3"
+
+# isort: off
+
+from lamini.api.lamini import Lamini
+from lamini.api.classifier import Classifier
+from lamini.api.embedding import Embedding
+from lamini.api.model_downloader import ModelDownloader
+from lamini.api.model_downloader import ModelType
+from lamini.api.model_downloader import DownloadedModel
+from lamini.classify.lamini_classifier import LaminiClassifier
+from lamini.generation.generation_node import GenerationNode
+from lamini.generation.generation_pipeline import GenerationPipeline
+from lamini.generation.base_prompt_object import PromptObject
+from lamini.generation.split_response_node import SplitResponseNode
+from lamini.api.streaming_completion import StreamingCompletion

lamini/api/classifier.py

@@ -1,5 +1,4 @@
-import time
-from typing import List, Union, Optional
+from typing import List, Optional, Union
 
 import lamini
 import requests

lamini/api/lamini.py

@@ -1,19 +1,20 @@
+import enum
 import json
-import jsonlines
 import logging
 import os
-import pandas as pd
 import time
+from typing import Any, Dict, Generator, Iterable, List, Optional, Union
 
-from lamini.api.lamini_config import get_config
-from lamini.api.rest_requests import get_version
+import jsonlines
+import pandas as pd
+import lamini
+from lamini.api.lamini_config import get_config, get_configured_key, get_configured_url
+from lamini.api.model_downloader import ModelDownloader, ModelType, DownloadedModel
+from lamini.api.rest_requests import get_version, make_web_request
 from lamini.api.train import Train
 from lamini.api.utils.completion import Completion
 from lamini.api.utils.upload_client import upload_to_blob
-from lamini.error.error import (
-    DownloadingModelError,
-)
-from typing import Dict, Iterable, List, Optional, Union, Any, Generator
+from lamini.error.error import DownloadingModelError
 
 logger = logging.getLogger(__name__)
 
@@ -48,29 +49,32 @@ def __init__(
         model_name: str,
         api_key: Optional[str] = None,
         api_url: Optional[str] = None,
+        model_type: ModelType = ModelType.transformer,
     ):
         self.config = get_config()
+        api_key = api_key or lamini.api_key or get_configured_key(self.config)
+        api_url = api_url or lamini.api_url or get_configured_url(self.config)
+
         self.model_name = model_name
         self.api_key = api_key
         self.api_url = api_url
         self.completion = Completion(api_key, api_url)
         self.trainer = Train(api_key, api_url)
         self.upload_file_path = None
         self.upload_base_path = None
+        self.model_downloader = ModelDownloader(api_key, api_url)
+        self.model_type = model_type
 
     def version(self) -> str:
         """Get the version of the Lamini platform
-
         Parameters
         ----------
         None
-
         Returns
         -------
         str
             Returned version fo the platform
         """
-
         return get_version(self.api_key, self.api_url, self.config)
 
     def generate(
@@ -117,17 +121,13 @@ def generate(
             specified, otherwise a dictionary matching the output_type is returned.
         """
 
-        result = None
-        try:
-            result = self.completion.generate(
-                prompt=prompt,
-                model_name=model_name or self.model_name,
-                output_type=output_type,
-                max_tokens=max_tokens,
-                max_new_tokens=max_new_tokens,
-            )
-        except DownloadingModelError as e:
-            return e
+        result = self.completion.generate(
+            prompt=prompt,
+            model_name=model_name or self.model_name,
+            output_type=output_type,
+            max_tokens=max_tokens,
+            max_new_tokens=max_new_tokens,
+        )
         if output_type is None:
             if isinstance(prompt, list):
                 result = [single_result["output"] for single_result in result]
@@ -370,6 +370,63 @@ def _upload_file_impl(
             )
         return items
 
+    def download_model(
+        self,
+        model_name: Optional[str] = None,
+        model_type: Optional[ModelType] = None,
+        wait: bool = False,
+        wait_time_seconds: int = 60,
+    ) -> DownloadedModel:
+        """Request Lamini Platform to download and cache the specified hugging face model.
+        So that the model can be immediately loaded to GPU memory afterwards.
+        Right now, only support downloading models from Hugging Face.
+
+        Parameters
+        ----------
+        hf_model_name: str
+            The full name of a hugging face model. Like meta-llama/Llama-3.2-11B-Vision-Instruct
+            in https://huggingface.co/meta-llama/Llama-3.2-11B-Vision-Instruct
+
+        model_type: ModelType
+            The type of the requested model.
+
+        Raises
+        ------
+        Exception
+            Raised if there is an issue with upload
+
+        Returns
+        -------
+        DownloadedModel
+        """
+        model_name_to_download = self.model_name if model_name is None else model_name
+        model_type_to_download = self.model_type if model_type is None else model_type
+        if not wait:
+            return self.model_downloader.download(
+                model_name_to_download, model_type_to_download
+            )
+
+        start_time = time.time()
+
+        while True:
+            result = self.model_downloader.download(
+                model_name_to_download, model_type_to_download
+            )
+
+            # Check the status of foo()'s result
+            if result.status == "available":
+                return result
+
+            # Check if the specified timeout has been exceeded
+            elapsed_time = time.time() - start_time
+            if elapsed_time > wait_time_seconds:
+                return result
+            INTERVAL_SECONDS = 1
+            time.sleep(INTERVAL_SECONDS)
+
+    def list_models(self) -> List[DownloadedModel]:
+        return self.model_downloader.list()
+
     def train(
         self,
         data_or_dataset_id: Union[
@@ -378,6 +435,7 @@ def train(
         finetune_args: Optional[dict] = None,
         gpu_config: Optional[dict] = None,
         is_public: Optional[bool] = None,
+        custom_model_name: Optional[str] = None,
     ) -> str:
         """Handler for training jobs through the Trainer object. This submits a training
         job request to the platform using the provided data.
@@ -398,6 +456,9 @@ def train(
         is_public: Optional[bool] = None
             Allow public access to the model and dataset
 
+        custom_model_name: Optional[str] = None
+            A human-readable name for the model.
+
         Raises
         ------
         AssertionError
@@ -429,6 +490,7 @@ def train(
             finetune_args=finetune_args,
             gpu_config=gpu_config,
             is_public=is_public,
+            custom_model_name=custom_model_name,
         )
         job["dataset_id"] = dataset_id
         return job
@@ -445,6 +507,7 @@ def train_and_wait(
         finetune_args: Optional[dict] = None,
         gpu_config: Optional[dict] = None,
         is_public: Optional[bool] = None,
+        custom_model_name: Optional[str] = None,
         **kwargs,
     ) -> str:
         """Handler for training jobs through the Trainer object. This submits a training
@@ -467,6 +530,9 @@ def train_and_wait(
         is_public: Optional[bool] = None
             Allow public access to the model and dataset
 
+        custom_model_name: Optional[str] = None
+            A human-readable name for the model.
+
         kwargs: Dict[str, Any]
             Key word arguments
                 verbose
@@ -488,6 +554,7 @@ def train_and_wait(
             finetune_args=finetune_args,
             gpu_config=gpu_config,
             is_public=is_public,
+            custom_model_name=custom_model_name,
         )
 
         try:

lamini/api/lamini_config.py

@@ -176,13 +176,21 @@ def get_configured_url(config: config.Configuration) -> str:
         Extracted platform url
     """
 
-    environment = os.environ.get("LLAMA_ENVIRONMENT")
+    LAMINI_ENV_ENV_VAR_NAME = "LLAMA_ENVIRONMENT"
+    LOCAL_URL_KEY = "local.url"
+    LOCAL_URL_DEFAULT_VALUE = "http://localhost:5001"
+    STAGING_URL_KEY = "staging.url"
+    STAGING_URL_DEFAULT_VALUE = "https://staging.lamini.ai"
+    PUBLIC_URL_DEFAULT_VALUE = "https://api.lamini.ai"
+    PUBLIC_URL_KEY = "production.url"
+
+    environment = os.environ.get(LAMINI_ENV_ENV_VAR_NAME)
     if environment == "LOCAL":
-        url = config.get("local.url", "http://localhost:5001")
+        url = config.get(LOCAL_URL_KEY, LOCAL_URL_DEFAULT_VALUE)
     elif environment == "STAGING":
-        url = config.get("staging.url", "https://staging.lamini.ai")
+        url = config.get(STAGING_URL_KEY, STAGING_URL_DEFAULT_VALUE)
     else:
-        url = config.get("production.url", "https://api.lamini.ai")
+        url = config.get(PUBLIC_URL_KEY, PUBLIC_URL_DEFAULT_VALUE)
     return url
 
 

lamini/api/model_downloader.py

@@ -0,0 +1,95 @@
+import enum
+from typing import List, Union
+
+import lamini
+import numpy as np
+from lamini.api.lamini_config import get_config, get_configured_key, get_configured_url
+from lamini.api.rest_requests import make_web_request
+
+
+class DownloadedModel:
+
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
+
+    model_id = None
+    model_name = None
+    model_type = None
+    user_id = None
+    is_public = None
+    creation_ts = None
+    prev_download_ts = None
+    prev_download_error = None
+    download_attempts = None
+    status = None
+
+    def __repr__(self):
+        return f"<DownloadedModel({', '.join(f'{k}={v!r}' for k, v in self.__dict__.items())})>"
+
+
+class ModelType(enum.Enum):
+    """This must be consistent with the db/migrations table definition's MODEL_TYPE type."""
+
+    transformer = "transformer"
+    embedding = "embedding"
+
+
+class ModelDownloader:
+    """Handler for requesting Lamini Platform to download a hugging face model.
+
+    Parameters
+    ----------
+    api_key: Optional[str]
+        Lamini platform API key, if not provided the key stored
+        within ~.lamini/configure.yaml will be used. If either
+        don't exist then an error is raised.
+
+    api_url: Optional[str]
+        Lamini platform api url, only needed if a different url is needed outside of the
+        defined ones here: https://github.com/lamini-ai/lamini-platform/blob/main/sdk/lamini/api/lamini_config.py#L68
+            i.e. localhost, staging.lamini.ai, or api.lamini.ai
+            Additionally, LLAMA_ENVIRONMENT can be set as an environment variable
+            that will be grabbed for the url before any of the above defaults
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        api_url: str,
+    ):
+        self.api_key = api_key
+        self.api_endpoint = api_url + "/v1/downloaded_models/"
+
+    def download(self, hf_model_name: str, model_type: ModelType) -> DownloadedModel:
+        """Request to Lamini platform for an embedding encoding of the provided
+        prompt
+
+
+        Parameters
+        ----------
+        prompt: Union[str, List[str]]
+            Prompt to encoding into an embedding
+
+        Returns
+        -------
+        DownloadedModel:
+            A object describing the state of the model.
+        """
+
+        params = {"hf_model_name": hf_model_name, "model_type": model_type.value}
+        resp = make_web_request(self.api_key, self.api_endpoint, "post", params)
+        return DownloadedModel(**resp)
+
+    def list(self) -> List[DownloadedModel]:
+        """List all models on the Lamini Platform.
+
+        Returns
+        -------
+        List[DownloadedModel]:
+            A object describing the state of the model.
+        """
+        resp = make_web_request(self.api_key, self.api_endpoint, "get")
+        res = []
+        for model in resp:
+            res.append(DownloadedModel(**model))
+        return res