lm-deluge 0.0.24__py3-none-any.whl → 0.0.26__py3-none-any.whl

lm_deluge/batches.py

@@ -169,7 +169,8 @@ async def submit_batches_oa(
     batch_size: int = 50_000,
 ):
     """Write OpenAI batch requests to a file and submit."""
-    BATCH_SIZE = batch_size
+    MAX_BATCH_SIZE_BYTES = 200 * 1024 * 1024  # 200MB
+    MAX_BATCH_SIZE_ITEMS = batch_size
 
     prompts = prompts_to_conversations(prompts)
     if any(p is None for p in prompts):
@@ -178,29 +179,71 @@ async def submit_batches_oa(
     model_obj = APIModel.from_registry(model)
 
     tasks = []
+    current_batch = []
+    current_batch_size = 0
+    # current_batch_start_idx = 0
+
+    for idx, prompt in enumerate(prompts):
+        assert isinstance(prompt, Conversation)
+        context = RequestContext(
+            task_id=idx,
+            model_name=model,
+            prompt=prompt,
+            sampling_params=sampling_params,
+        )
+        request = {
+            "custom_id": str(idx),
+            "method": "POST",
+            "url": "/v1/chat/completions",
+            "body": await _build_oa_chat_request(model_obj, context),
+        }
 
-    for start in range(0, len(prompts), BATCH_SIZE):
-        batch_prompts = prompts[start : start + BATCH_SIZE]
-        with tempfile.NamedTemporaryFile(mode="w+", suffix=".jsonl", delete=False) as f:
-            for idx, prompt in enumerate(batch_prompts, start=start):
-                assert isinstance(prompt, Conversation)
-                context = RequestContext(
-                    task_id=idx,
-                    model_name=model,
-                    prompt=prompt,
-                    sampling_params=sampling_params,
-                )
-                request = {
-                    "custom_id": str(idx),
-                    "method": "POST",
-                    "url": "/v1/chat/completions",
-                    "body": await _build_oa_chat_request(model_obj, context),
-                }
-                json.dump(request, f)
-                f.write("\n")
-
-            file_path = f.name
-
+        # Calculate size of this request
+        request_json = json.dumps(request) + "\n"
+        request_size = len(request_json.encode("utf-8"))
+
+        # Check if adding this request would exceed limits
+        would_exceed_size = current_batch_size + request_size > MAX_BATCH_SIZE_BYTES
+        would_exceed_items = len(current_batch) >= MAX_BATCH_SIZE_ITEMS
+
+        if current_batch and (would_exceed_size or would_exceed_items):
+            # Submit current batch
+            def write_batch_file():
+                with tempfile.NamedTemporaryFile(
+                    mode="w+", suffix=".jsonl", delete=False
+                ) as f:
+                    for batch_request in current_batch:
+                        json.dump(batch_request, f)
+                        f.write("\n")
+                    print("wrote", len(current_batch), "items")
+                    return f.name
+
+            file_path = await asyncio.to_thread(write_batch_file)
+            tasks.append(asyncio.create_task(submit_batch_oa(file_path)))
+
+            # Start new batch
+            current_batch = []
+            current_batch_size = 0
+            # current_batch_start_idx = idx
+
+        # Add request to current batch
+        current_batch.append(request)
+        current_batch_size += request_size
+
+    # Submit final batch if it has items
+    if current_batch:
+
+        def write_final_batch_file():
+            with tempfile.NamedTemporaryFile(
+                mode="w+", suffix=".jsonl", delete=False
+            ) as f:
+                for batch_request in current_batch:
+                    json.dump(batch_request, f)
+                    f.write("\n")
+                print("wrote", len(current_batch), "items")
+                return f.name
+
+        file_path = await asyncio.to_thread(write_final_batch_file)
         tasks.append(asyncio.create_task(submit_batch_oa(file_path)))
 
     batch_ids = await asyncio.gather(*tasks)
@@ -229,34 +272,80 @@ async def submit_batches_anthropic(
 
     Returns: batch_ids (list[str])
     """
+    MAX_BATCH_SIZE_BYTES = 200 * 1024 * 1024  # 200MB
+    MAX_BATCH_SIZE_ITEMS = batch_size
 
     # Convert prompts to Conversations
     prompts = prompts_to_conversations(prompts)
 
     request_headers = None
-    BATCH_SIZE = batch_size
     batch_tasks = []
-
-    for start in range(0, len(prompts), BATCH_SIZE):
-        batch_prompts = prompts[start : start + BATCH_SIZE]
-        with tempfile.NamedTemporaryFile(mode="w+", suffix=".jsonl", delete=False) as f:
-            for idx, prompt in enumerate(batch_prompts, start=start):
-                assert isinstance(prompt, Conversation)
-                context = RequestContext(
-                    task_id=idx,
-                    model_name=model,
-                    prompt=prompt,
-                    sampling_params=sampling_params,
-                    cache=cache,
-                )
-                request_body, request_headers = _build_anthropic_request(
-                    APIModel.from_registry(model), context
+    current_batch = []
+    current_batch_size = 0
+
+    for idx, prompt in enumerate(prompts):
+        assert isinstance(prompt, Conversation)
+        context = RequestContext(
+            task_id=idx,
+            model_name=model,
+            prompt=prompt,
+            sampling_params=sampling_params,
+            cache=cache,
+        )
+        request_body, request_headers = _build_anthropic_request(
+            APIModel.from_registry(model), context
+        )
+        request = {"custom_id": str(idx), "params": request_body}
+
+        # Calculate size of this request
+        request_json = json.dumps(request) + "\n"
+        request_size = len(request_json.encode("utf-8"))
+
+        # Check if adding this request would exceed limits
+        would_exceed_size = current_batch_size + request_size > MAX_BATCH_SIZE_BYTES
+        would_exceed_items = len(current_batch) >= MAX_BATCH_SIZE_ITEMS
+
+        if current_batch and (would_exceed_size or would_exceed_items):
+            # Submit current batch
+            def write_batch_file():
+                with tempfile.NamedTemporaryFile(
+                    mode="w+", suffix=".jsonl", delete=False
+                ) as f:
+                    for batch_request in current_batch:
+                        json.dump(batch_request, f)
+                        f.write("\n")
+                    print("wrote", len(current_batch), "items")
+                    return f.name
+
+            file_path = await asyncio.to_thread(write_batch_file)
+            batch_tasks.append(
+                asyncio.create_task(
+                    _submit_anthropic_batch(file_path, request_headers, model)  # type: ignore
                 )
-                json.dump({"custom_id": str(idx), "params": request_body}, f)
-                f.write("\n")
+            )
+
+            # Start new batch
+            current_batch = []
+            current_batch_size = 0
+
+        # Add request to current batch
+        current_batch.append(request)
+        current_batch_size += request_size
+
+    # Submit final batch if it has items
+    if current_batch:
 
-            file_path = f.name
+        def write_final_batch_file():
+            with tempfile.NamedTemporaryFile(
+                mode="w+", suffix=".jsonl", delete=False
+            ) as f:
+                for batch_request in current_batch:
+                    json.dump(batch_request, f)
+                    f.write("\n")
+                print("wrote", len(current_batch), "items")
+                return f.name
 
+        file_path = await asyncio.to_thread(write_final_batch_file)
         batch_tasks.append(
             asyncio.create_task(
                 _submit_anthropic_batch(file_path, request_headers, model)  # type: ignore

lm_deluge/file.py

@@ -1,3 +1,4 @@
+from functools import cached_property
 import os
 import io
 import requests
@@ -68,13 +69,13 @@ class File:
             return encoded
         return f"data:{self._mime()};base64,{encoded}"
 
-    @property
+    @cached_property
     def fingerprint(self) -> str:
         # Hash the file contents for fingerprinting
         file_bytes = self._bytes()
         return xxhash.xxh64(file_bytes).hexdigest()
 
-    @property
+    @cached_property
     def size(self) -> int:
         """Return file size in bytes."""
         return len(self._bytes())

lm_deluge/request_context.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass, field
+from functools import cached_property
 from typing import Any, Callable
 
 from .config import SamplingParams
@@ -39,14 +40,18 @@ class RequestContext:
 
     # Computed properties
     cache_key: str = field(init=False)
-    num_tokens: int = field(init=False)
+    # num_tokens: int = field(init=False)
 
-    def __post_init__(self):
-        # Compute cache key from prompt fingerprint
-        self.cache_key = self.prompt.fingerprint
+    # def __post_init__(self):
+    #     # Compute cache key from prompt fingerprint
+    #     # self.cache_key = self.prompt.fingerprint
 
-        # Compute token count
-        self.num_tokens = self.prompt.count_tokens(self.sampling_params.max_new_tokens)
+    #     # Compute token count
+    #     self.num_tokens =
+
+    @cached_property
+    def num_tokens(self):
+        return self.prompt.count_tokens(self.sampling_params.max_new_tokens)
 
     def maybe_callback(self, response, tracker):
         if not self.callback:

{lm_deluge-0.0.24.dist-info → lm_deluge-0.0.26.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lm_deluge
-Version: 0.0.24
+Version: 0.0.26
 Summary: Python utility for using LLM API models.
 Author-email: Benjamin Anderson <ben@trytaylor.ai>
 Requires-Python: >=3.10

{lm_deluge-0.0.24.dist-info → lm_deluge-0.0.26.dist-info}/RECORD

@@ -1,17 +1,17 @@
 lm_deluge/__init__.py,sha256=mAztMuxINmh7dGbYnT8tsmw1eryQAvd0jpY8yHzd0EE,315
 lm_deluge/agent.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-lm_deluge/batches.py,sha256=Q7uETobLxgMJ7_Y_oJU1MKOYCa7ttmCU6CX6EuPv2qY,18636
+lm_deluge/batches.py,sha256=vJXVnuuGkIQnXoDPODPERrvdG9X1Ov1jnXExnPe6ZAc,21772
 lm_deluge/cache.py,sha256=VB1kv8rM2t5XWPR60uhszFcxLDnVKOe1oA5hYjVDjIo,4375
 lm_deluge/client.py,sha256=nAGMwdUPDVx-x23hZF6U5Yhug6Zf5FT27RHj_kj8nZk,26369
 lm_deluge/config.py,sha256=H1tQyJDNHGFuwxqQNL5Z-CjWAC0luHSBA3iY_pxmACM,932
 lm_deluge/embed.py,sha256=CO-TOlC5kOTAM8lcnicoG4u4K664vCBwHF1vHa-nAGg,13382
 lm_deluge/errors.py,sha256=oHjt7YnxWbh-eXMScIzov4NvpJMo0-2r5J6Wh5DQ1tk,209
-lm_deluge/file.py,sha256=zQH1STMjCG9pczO7Fk9Jw0_0Pj_8CogcdIxTe4J4AJw,5414
+lm_deluge/file.py,sha256=ggOXcAOI-LAykXLb3N3WfLlHdBjQIJtW9ltmfKCIsNQ,5466
 lm_deluge/gemini_limits.py,sha256=V9mpS9JtXYz7AY6OuKyQp5TuIMRH1BVv9YrSNmGmHNA,1569
 lm_deluge/image.py,sha256=Qpa0k5yXfrpSaHzVUwW_TEn7yEgmwzYGL17Sa7-KhSA,7729
 lm_deluge/models.py,sha256=3vgI1BlfT4_Higev25QhhXJufQvsI6pd0yjF9YL0crA,49812
 lm_deluge/prompt.py,sha256=cfwzCAmT-1K0v7SfEMUrxpBkJGgf7IFlWfNLJrCcoBM,37025
-lm_deluge/request_context.py,sha256=0X-5m8BKn51rnnjzGDDXqbuSUEFGjdayirQjbvPcjMI,2425
+lm_deluge/request_context.py,sha256=o33LSEwnK6YPhZeulUoSE_VrdKCXiCQa0tjjixK2K6M,2540
 lm_deluge/rerank.py,sha256=-NBAJdHz9OB-SWWJnHzkFmeVO4wR6lFV7Vw-SxG7aVo,11457
 lm_deluge/tool.py,sha256=_coOKB9nPNVZoseMRumRyQ8BMR7_d0IlstzMHNT69JY,15732
 lm_deluge/tracker.py,sha256=-EkFDAklh5mclIFR-5SthAwNL4p1yKS8LUN7rhpOVPQ,9266
@@ -48,8 +48,8 @@ lm_deluge/util/logprobs.py,sha256=UkBZakOxWluaLqHrjARu7xnJ0uCHVfLGHJdnYlEcutk,11
 lm_deluge/util/spatial.py,sha256=BsF_UKhE-x0xBirc-bV1xSKZRTUhsOBdGqsMKme20C8,4099
 lm_deluge/util/validation.py,sha256=hz5dDb3ebvZrZhnaWxOxbNSVMI6nmaOODBkk0htAUhs,1575
 lm_deluge/util/xml.py,sha256=Ft4zajoYBJR3HHCt2oHwGfymGLdvp_gegVmJ-Wqk4Ck,10547
-lm_deluge-0.0.24.dist-info/licenses/LICENSE,sha256=uNNXGXPCw2TC7CUs7SEBkA-Mz6QBQFWUUEWDMgEs1dU,1058
-lm_deluge-0.0.24.dist-info/METADATA,sha256=mXZ-AgyVqmY5fD_U8qbCb0Hm5Ekx4jkSyW8dQ3zYDJw,12978
-lm_deluge-0.0.24.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-lm_deluge-0.0.24.dist-info/top_level.txt,sha256=hqU-TJX93yBwpgkDtYcXyLr3t7TLSCCZ_reytJjwBaE,10
-lm_deluge-0.0.24.dist-info/RECORD,,
+lm_deluge-0.0.26.dist-info/licenses/LICENSE,sha256=uNNXGXPCw2TC7CUs7SEBkA-Mz6QBQFWUUEWDMgEs1dU,1058
+lm_deluge-0.0.26.dist-info/METADATA,sha256=2zzPvIBE0SdA0RAAKGnyUQDIq-sE_Eo34pfRu_ygi4I,12978
+lm_deluge-0.0.26.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+lm_deluge-0.0.26.dist-info/top_level.txt,sha256=hqU-TJX93yBwpgkDtYcXyLr3t7TLSCCZ_reytJjwBaE,10
+lm_deluge-0.0.26.dist-info/RECORD,,