Update to 3.0.4 (#67)

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

lamini/api/utils/completion.py

@@ -74,7 +74,7 @@ def make_llm_req_map(
         # TODO: prompt should be named prompt to signal it's a batch.
         if isinstance(prompt, list) and len(prompt) > 20:
             print(
-                "For large inference batches, consider using a Generation Pipeline instead: https://github.com/lamini-ai/lamini-examples/blob/main/05_data_pipeline/README.md"
+                "For large inference batches, we strongly recommend using a Generation Pipeline to streamline your process: https://github.com/lamini-ai/lamini-examples/blob/main/05_data_pipeline/"
             )
         req_data["prompt"] = prompt
         req_data["output_type"] = output_type

lamini/api/utils/reservations.py

@@ -110,28 +110,9 @@ async def wait_and_poll_for_reservation(self, client):
         async with self.condition:
             self.condition.notify(len(self.condition._waiters))
         self.is_polling = False
-        if self.is_working:
-            self.polling_task = asyncio.create_task(
-                self.kickoff_reservation_polling(client)
-            )
-            logger.info("Made reservation " + str(reservation))
-            if "dynamic_max_batch_size" not in reservation:
-                reservation["dynamic_max_batch_size"] = lamini.batch_size
-            self.current_reservation = reservation
-            self.capacity_remaining = reservation["capacity_remaining"]
-            self.dynamic_max_batch_size = reservation["dynamic_max_batch_size"]
-            if self.variable_capacity:
-                self.capacity_needed = self.dynamic_max_batch_size * lamini.max_workers
-            async with self.condition:
-                self.condition.notify(len(self.condition._waiters))
-            self.is_polling = False
-            if self.is_working:
-                self.polling_task = asyncio.create_task(
-                    self.kickoff_reservation_polling(client)
-                )
-                _ = asyncio.create_task(
-                    self.timer_based_polling(reservation["end_time"])
-                )
+        self.polling_task = asyncio.create_task(
+            self.kickoff_reservation_polling(client)
+        )
 
     async def timer_based_polling(self, wakeup_time):
         try:

lamini/classify/lamini_classifier.py

@@ -62,30 +62,30 @@ def prompt_train(self, prompts: dict):
 
         First, augment the examples for each class using the prompts.
         """
-        for class_name, prompt in prompts.items():
-            try:
-                logger.info(
-                    f"Generating examples for class '{class_name}' from prompt {prompt}"
-                )
-                self.add_class(class_name)
+        try:
+            for class_name, prompt in prompts.items():
+                    logger.info(
+                        f"Generating examples for class '{class_name}' from prompt {prompt}"
+                    )
+                    self.add_class(class_name)
 
-                result = self.generate_examples_from_prompt(
-                    class_name, prompt, self.examples.get(class_name, [])
-                )
+                    result = self.generate_examples_from_prompt(
+                        class_name, prompt, self.examples.get(class_name, [])
+                    )
 
-                self.examples[class_name] = result
+                    self.examples[class_name] = result
 
-                # Save partial progress
-                self.save_examples()
+                    # Save partial progress
+                    self.save_examples()
 
-            except Exception as e:
-                logger.error(f"Failed to generate examples for class {class_name}")
-                logger.error(e)
-                logger.error(
-                    "Consider rerunning the generation task if the error is transient, e.g. 500"
-                )
 
-        self.train()
+            self.train()
+        except Exception as e:
+            logger.error(f"Failed to generate examples for class {class_name}")
+            logger.error(e)
+            logger.error(
+                "Consider rerunning the generation task if the error is transient, e.g. 500"
+            )
 
     def train(self):
         # Form the embeddings

lamini/generation/classifier_node.py

@@ -14,13 +14,10 @@ def __init__(
         self,
         classifier: LaminiClassifier,
         model_name: Optional[str] = None,
-        api_key: Optional[str] = None,
-        api_url: Optional[str] = None,
-        config: dict = {},
         max_tokens: Optional[int] = None,
     ):
         super(ClassifierNode, self).__init__(
-            model_name=model_name, api_key=api_key, api_url=api_url, config=config
+            model_name=model_name
         )
         self.max_tokens = max_tokens
         self.classifier = classifier

lamini/generation/embedding_node.py

@@ -2,7 +2,6 @@
 import sys
 from typing import AsyncIterator, Iterator, Optional, Union
 
-from lamini.api.lamini_config import get_config
 from lamini.generation.base_prompt_object import PromptObject
 from lamini.generation.generation_node import GenerationNode
 
@@ -28,50 +27,14 @@ class EmbeddingNode(GenerationNode):
     ----------
     model_name: Optional[str]
         Model name as referred to on HuggingFace https://huggingface.co/models
-
-    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
-
-    config: dict
-        Dictionary that is handled from the following script:
-            https://github.com/lamini-ai/lamini-platform/blob/main/sdk/lamini/api/lamini_config.py
-        Configurations currently hold the following keys and data as a yaml format:
-            local:
-                url: <url>
-            staging:
-                url: <url>
-            production:
-                url: <url>
-
-            local:
-                key: <auth-key>
-            staging:
-                key: <auth-key>
-            production:
-                key:
-                    <auth-key
-
-
     """
 
     def __init__(
         self,
         model_name: Optional[str] = None,
-        api_key: Optional[str] = None,
-        api_url: Optional[str] = None,
-        config: dict = {},
     ):
         super(EmbeddingNode, self).__init__(
-            model_name=model_name, api_key=api_key, api_url=api_url, config=config
+            model_name=model_name
         )
 
     def generate(

lamini/generation/index_node.py

@@ -14,13 +14,10 @@ def __init__(
         index: LaminiIndex,
         index_top_k: int = 1,
         model_name: Optional[str] = None,
-        api_key: Optional[str] = None,
-        api_url: Optional[str] = None,
-        config: dict = {},
         max_tokens: Optional[int] = None,
     ):
         super(IndexNode, self).__init__(
-            model_name=model_name, api_key=api_key, api_url=api_url, config=config
+            model_name=model_name
         )
         self.max_tokens = max_tokens
         self.index = index

lamini/generation/process_generation_batch.py

@@ -19,6 +19,7 @@ def can_submit_query():
         if reservation_api.current_reservation is None:
             return True
         if reservation_api.capacity_remaining < len(batch["prompt"]):
+            reservation_api.poll_for_reservation.set()
             return False
         # Now we can consume credits and send batch
         reservation_api.update_capacity_use(len(batch["prompt"]))

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "lamini"
-version = "3.0.0"
+version = "3.0.4"
 authors = [
   { name="PowerML", email="[email protected]" },
 ]