Merge branch 'main' of github.com:expectedparrot/edsl into nb_deletes

edsl/jobs/buckets/BucketCollection.py

@@ -13,6 +13,8 @@ class BucketCollection(UserDict):
     def __init__(self, infinity_buckets=False):
         super().__init__()
         self.infinity_buckets = infinity_buckets
+        self.models_to_services = {}
+        self.services_to_buckets = {}
 
     def __repr__(self):
         return f"BucketCollection({self.data})"
@@ -21,6 +23,7 @@ def add_model(self, model: "LanguageModel") -> None:
         """Adds a model to the bucket collection.
 
         This will create the token and request buckets for the model."""
+
         # compute the TPS and RPS from the model
         if not self.infinity_buckets:
             TPS = model.TPM / 60.0
@@ -29,22 +32,28 @@ def add_model(self, model: "LanguageModel") -> None:
             TPS = float("inf")
             RPS = float("inf")
 
-        # create the buckets
-        requests_bucket = TokenBucket(
-            bucket_name=model.model,
-            bucket_type="requests",
-            capacity=RPS,
-            refill_rate=RPS,
-        )
-        tokens_bucket = TokenBucket(
-            bucket_name=model.model, bucket_type="tokens", capacity=TPS, refill_rate=TPS
-        )
-        model_buckets = ModelBuckets(requests_bucket, tokens_bucket)
-        if model in self:
-            # it if already exists, combine the buckets
-            self[model] += model_buckets
+        if model.model not in self.models_to_services:
+            service = model._inference_service_
+            if service not in self.services_to_buckets:
+                requests_bucket = TokenBucket(
+                    bucket_name=service,
+                    bucket_type="requests",
+                    capacity=RPS,
+                    refill_rate=RPS,
+                )
+                tokens_bucket = TokenBucket(
+                    bucket_name=service,
+                    bucket_type="tokens",
+                    capacity=TPS,
+                    refill_rate=TPS,
+                )
+                self.services_to_buckets[service] = ModelBuckets(
+                    requests_bucket, tokens_bucket
+                )
+            self.models_to_services[model.model] = service
+            self[model] = self.services_to_buckets[service]
         else:
-            self[model] = model_buckets
+            self[model] = self.services_to_buckets[self.models_to_services[model.model]]
 
     def visualize(self) -> dict:
         """Visualize the token and request buckets for each model."""

edsl/language_models/LanguageModel.py

@@ -164,10 +164,6 @@ class LanguageModel(
         None  # This should be something like ["choices", 0, "message", "content"]
     )
     __rate_limits = None
-    __default_rate_limits = {
-        "rpm": 10_000,
-        "tpm": 2_000_000,
-    }  # TODO: Use the OpenAI Teir 1 rate limits
     _safety_factor = 0.8
 
     def __init__(
@@ -181,6 +177,7 @@ def __init__(
         self.remote = False
         self.omit_system_prompt_if_empty = omit_system_prompt_if_empty_string
 
+        # self._rpm / _tpm comes from the class
         if rpm is not None:
             self._rpm = rpm
 
@@ -289,35 +286,40 @@ def set_rate_limits(self, rpm=None, tpm=None) -> None:
         >>> m.RPM
         100
         """
-        self._set_rate_limits(rpm=rpm, tpm=tpm)
+        if rpm is not None:
+            self._rpm = rpm
+        if tpm is not None:
+            self._tpm = tpm
+        return None
+        # self._set_rate_limits(rpm=rpm, tpm=tpm)
 
-    def _set_rate_limits(self, rpm=None, tpm=None) -> None:
-        """Set the rate limits for the model.
+    # def _set_rate_limits(self, rpm=None, tpm=None) -> None:
+    #     """Set the rate limits for the model.
 
-        If the model does not have rate limits, use the default rate limits."""
-        if rpm is not None and tpm is not None:
-            self.__rate_limits = {"rpm": rpm, "tpm": tpm}
-            return
+    #     If the model does not have rate limits, use the default rate limits."""
+    #     if rpm is not None and tpm is not None:
+    #         self.__rate_limits = {"rpm": rpm, "tpm": tpm}
+    #         return
 
-        if self.__rate_limits is None:
-            if hasattr(self, "get_rate_limits"):
-                self.__rate_limits = self.get_rate_limits()
-            else:
-                self.__rate_limits = self.__default_rate_limits
+    #     if self.__rate_limits is None:
+    #         if hasattr(self, "get_rate_limits"):
+    #             self.__rate_limits = self.get_rate_limits()
+    #         else:
+    #             self.__rate_limits = self.__default_rate_limits
 
     @property
     def RPM(self):
         """Model's requests-per-minute limit."""
         # self._set_rate_limits()
         # return self._safety_factor * self.__rate_limits["rpm"]
-        return self.rpm
+        return self._rpm
 
     @property
     def TPM(self):
         """Model's tokens-per-minute limit."""
         # self._set_rate_limits()
         # return self._safety_factor * self.__rate_limits["tpm"]
-        return self.tpm
+        return self._tpm
 
     @property
     def rpm(self):
@@ -335,17 +337,6 @@ def tpm(self):
     def tpm(self, value):
         self._tpm = value
 
-    @property
-    def TPM(self):
-        """Model's tokens-per-minute limit.
-
-        >>> m = LanguageModel.example()
-        >>> m.TPM > 0
-        True
-        """
-        self._set_rate_limits()
-        return self._safety_factor * self.__rate_limits["tpm"]
-
     @staticmethod
     def _overide_default_parameters(passed_parameter_dict, default_parameter_dict):
         """Return a dictionary of parameters, with passed parameters taking precedence over defaults.

edsl/questions/derived/QuestionLinearScale.py

@@ -22,6 +22,7 @@ def __init__(
         option_labels: Optional[dict[int, str]] = None,
         answering_instructions: Optional[str] = None,
         question_presentation: Optional[str] = None,
+        include_comment: Optional[bool] = True,
     ):
         """Instantiate a new QuestionLinearScale.
 
@@ -36,6 +37,7 @@ def __init__(
             question_text=question_text,
             question_options=question_options,
             use_code=False,  # question linear scale will have it's own code
+            include_comment=include_comment,
         )
         self.question_options = question_options
         self.option_labels = (

edsl/questions/derived/QuestionTopK.py

@@ -20,6 +20,7 @@ def __init__(
         max_selections: int,
         question_presentation: Optional[str] = None,
         answering_instructions: Optional[str] = None,
+        include_comment: Optional[bool] = True,
     ):
         """Initialize the question.
 
@@ -37,6 +38,7 @@ def __init__(
             max_selections=max_selections,
             question_presentation=question_presentation,
             answering_instructions=answering_instructions,
+            include_comment=include_comment,
         )
         if min_selections != max_selections:
             raise QuestionCreationValidationError(

edsl/questions/derived/QuestionYesNo.py

@@ -19,6 +19,7 @@ def __init__(
         question_options: list[str] = ["No", "Yes"],
         answering_instructions: Optional[str] = None,
         question_presentation: Optional[str] = None,
+        include_comment: Optional[bool] = True,
     ):
         """Instantiate a new QuestionYesNo.
 
@@ -33,6 +34,7 @@ def __init__(
             use_code=False,
             answering_instructions=answering_instructions,
             question_presentation=question_presentation,
+            include_comment=include_comment,
         )
         self.question_options = question_options
 

edsl/scenarios/FileStore.py

@@ -241,8 +241,16 @@ def example(cls):
 
 
 class PNGFileStore(FileStore):
-    def __init__(self, filename):
-        super().__init__(filename, suffix=".png")
+    def __init__(
+        self,
+        filename,
+        binary: Optional[bool] = None,
+        suffix: Optional[str] = None,
+        base64_string: Optional[str] = None,
+    ):
+        super().__init__(
+            filename, binary=binary, base64_string=base64_string, suffix=".png"
+        )
 
     @classmethod
     def example(cls):
@@ -267,8 +275,17 @@ def view(self):
 
 
 class SQLiteFileStore(FileStore):
-    def __init__(self, filename):
-        super().__init__(filename, suffix=".sqlite")
+
+    def __init__(
+        self,
+        filename,
+        binary: Optional[bool] = None,
+        suffix: Optional[str] = None,
+        base64_string: Optional[str] = None,
+    ):
+        super().__init__(
+            filename, binary=binary, base64_string=base64_string, suffix=".sqlite"
+        )
 
     @classmethod
     def example(cls):
@@ -281,6 +298,8 @@ def example(cls):
             c.execute("""CREATE TABLE stocks (date text)""")
             conn.commit()
 
+            return cls(f.name)
+
     def view(self):
         import subprocess
         import os
@@ -290,8 +309,25 @@ def view(self):
 
 
 class HTMLFileStore(FileStore):
-    def __init__(self, filename):
-        super().__init__(filename, suffix=".html")
+
+    def __init__(
+        self,
+        filename,
+        binary: Optional[bool] = None,
+        suffix: Optional[str] = None,
+        base64_string: Optional[str] = None,
+    ):
+        super().__init__(
+            filename, binary=binary, base64_string=base64_string, suffix=".html"
+        )
+
+    @classmethod
+    def example(cls):
+        import tempfile
+
+        with tempfile.NamedTemporaryFile(suffix=".html", delete=False) as f:
+            f.write("<html><body><h1>Test</h1></body></html>".encode())
+            return cls(f.name)
 
     def view(self):
         import webbrowser

tests/base/test_Base.py

@@ -16,6 +16,7 @@ class SaveLoadFail(Warning):
 
 class TestBaseModels:
     def test_register_subclasses_meta(self):
+
         for key, value in RegisterSubclassesMeta.get_registry().items():
             assert key in [
                 "Result",
@@ -30,6 +31,12 @@ def test_register_subclasses_meta(self):
                 "Cache",
                 "Notebook",
                 "ModelList",
+                "FileStore",
+                "HTMLFileStore",
+                "CSVFileStore",
+                "PDFFileStore",
+                "PNGFileStore",
+                "SQLiteFileStore",
             ]
 
         methods = [

tests/jobs/test_BucketCollection.py

@@ -0,0 +1,14 @@
+import pytest
+
+from edsl import Model, QuestionFreeText
+
+
+def test_one_per_service():
+    models = [Model(temperature=1), Model(temperature=2), Model(temperature=0)]
+    q = QuestionFreeText(
+        question_text="What is your favorite color?", question_name="color"
+    )
+    jobs = q.by(models)
+    bc = jobs.bucket_collection
+    assert len(bc) == 3
+    assert len(set(bc.values())) == 1

tests/language_models/test_LanguageModel.py

@@ -21,6 +21,19 @@ class TestLanguageModel(unittest.TestCase):
     def setUp(self):
         pass
 
+    def test_tokens(self):
+        import random
+
+        random_tpm = random.randint(0, 100)
+        random_tpm = random.randint(0, 100)
+        m = LanguageModel.example()
+        m.set_rate_limits(tpm=random_tpm, rpm=random_tpm)
+        self.assertEqual(m.tpm, random_tpm)
+        self.assertEqual(m.rpm, random_tpm)
+
+        m.rpm = 45
+        self.assertEqual(m.rpm, 45)
+
     def test_execute_model_call(self):
         from edsl.data.Cache import Cache