lm-deluge 0.0.3__py3-none-any.whl

lm_deluge/embed.py

@@ -0,0 +1,392 @@
+### specific utility for cohere rerank api
+import os
+import numpy as np
+import aiohttp
+from tqdm.auto import tqdm
+import asyncio
+import time
+from typing import Any, Optional
+from dataclasses import dataclass
+from .tracker import StatusTracker
+
+registry = {
+    "text-embedding-3-small": {
+        "name": "text-embedding-3-small",
+        "provider": "openai",
+        "cost": 0.02,  # per million tokens
+    },
+    "text-embedding-3-large": {
+        "name": "text-embedding-3-large",
+        "provider": "openai",
+        "cost": 0.13,
+    },
+    "text-embedding-ada-002": {
+        "name": "text-embedding-ada-002",
+        "provider": "openai",
+        "cost": 1,
+    },
+    "embed-english-v3.0": {
+        "name": "embed-english-v3.0",
+        "provider": "cohere",
+        "cost": 0.1,
+    },
+    "embed-english-light-v3.0": {
+        "name": "embed-english-light-v3.0",
+        "provider": "cohere",
+        "cost": 0.1,
+    },
+    "embed-multilingual-v3.0": {
+        "name": "embed-multilingual-v3.0",
+        "provider": "cohere",
+        "cost": 0.1,
+    },
+    "embed-multilingual-light-v3.0": {
+        "name": "embed-multilingual-light-v3.0",
+        "provider": "cohere",
+        "cost": 0.1,
+    },
+}
+
+
+class EmbeddingRequest:
+    def __init__(
+        self,
+        task_id: int,
+        model_name: str,
+        texts: list[str],
+        attempts_left: int,
+        status_tracker: StatusTracker,
+        retry_queue: asyncio.Queue,
+        request_timeout: int,
+        pbar: Optional[tqdm] = None,
+        **kwargs,  # openai or cohere specific params
+    ):
+        self.task_id = task_id
+        self.model_name = model_name
+        self.texts = texts
+        self.attempts_left = attempts_left
+        self.status_tracker = status_tracker
+        self.retry_queue = retry_queue
+        self.request_timeout = request_timeout
+        self.pbar = pbar
+        self.result = []
+        self.kwargs = kwargs
+
+    def increment_pbar(self):
+        if self.pbar is not None:
+            self.pbar.update(1)
+
+    def handle_success(self):
+        self.increment_pbar()
+        self.status_tracker.num_tasks_in_progress -= 1
+        self.status_tracker.num_tasks_succeeded += 1
+
+    def handle_error(self):
+        last_result: EmbeddingResponse = self.result[-1]
+        error_to_print = (
+            f"Error on task {self.task_id}, Code: {last_result.status_code}, "
+        )
+        error_to_print += f"Message: {last_result.error_message}."
+        print(error_to_print)
+        if self.attempts_left > 0:
+            self.attempts_left -= 1
+            self.retry_queue.put_nowait(self)
+            return
+        else:
+            print(f"Task {self.task_id} out of tries.")
+            self.status_tracker.num_tasks_in_progress -= 1
+            self.status_tracker.num_tasks_failed += 1
+
+    async def handle_response(self, response: aiohttp.ClientResponse):
+        try:
+            if response.status == 200:
+                result = await response.json()
+                # TODO: add cost calculation
+                if self.model_name in [
+                    "text-embedding-3-small",
+                    "text-embedding-3-large",
+                    "text-embedding-ada-002",
+                ]:
+                    embeddings = [
+                        embedding["embedding"] for embedding in result["data"]
+                    ]
+                elif self.model_name in [
+                    "embed-english-v3.0",
+                    "embed-english-light-v3.0",
+                    "embed-multilingual-v3.0",
+                    "embed-multilingual-light-v3.0",
+                ]:
+                    embeddings = result["embeddings"]
+                else:
+                    raise ValueError(f"Unsupported model {self.model_name}")
+                return EmbeddingResponse(
+                    id=self.task_id,
+                    status_code=response.status,
+                    is_error=False,
+                    error_message=None,
+                    texts=self.texts,
+                    embeddings=embeddings,
+                )
+            else:
+                error_msg = await response.text()
+                return EmbeddingResponse(
+                    id=self.task_id,
+                    status_code=response.status,
+                    is_error=True,
+                    error_message=error_msg,
+                    texts=[],
+                    embeddings=[],
+                )
+        except Exception as e:
+            return EmbeddingResponse(
+                id=self.task_id,
+                status_code=response.status,
+                is_error=True,
+                error_message=str(e),
+                texts=[],
+                embeddings=[],
+            )
+
+    async def call_api(
+        self,
+        session: aiohttp.ClientSession,
+    ):
+        if len(self.texts) > 96:
+            raise ValueError("Embeddings only support up to 96 texts per request.")
+        model_obj = registry[self.model_name]
+        url = (
+            "https://api.openai.com/v1/embeddings"
+            if model_obj["provider"] == "openai"
+            else "https://api.cohere.com/v1/embed"
+        )
+        headers = {
+            "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"
+            if model_obj["provider"] == "openai"
+            else f"bearer {os.environ.get('COHERE_API_KEY')}"
+        }
+        payload: dict[str, Any] = {"model": self.model_name}
+        if model_obj["provider"] == "openai":
+            payload["input"] = self.texts
+            payload["encoding_format"] = "float"
+            for k, v in self.kwargs.items():
+                payload[k] = v
+        elif model_obj["provider"] == "cohere":
+            payload["texts"] = self.texts
+            payload["input_type"] = self.kwargs.get("input_type", "search_document")
+            for k, v in self.kwargs.items():
+                payload[k] = v
+        try:
+            self.status_tracker.total_requests += 1
+            async with session.post(url, json=payload, headers=headers) as response:
+                response_obj: EmbeddingResponse = await self.handle_response(response)
+            self.result.append(response_obj)
+            if response_obj.is_error:
+                self.handle_error()
+            else:
+                self.handle_success()
+
+        except asyncio.TimeoutError:
+            self.result.append(
+                EmbeddingResponse(
+                    id=self.task_id,
+                    status_code=None,
+                    is_error=True,
+                    error_message="Timeout",
+                    texts=[],
+                    embeddings=[],
+                )
+            )
+            self.handle_error()
+
+        except Exception as e:
+            self.result.append(
+                EmbeddingResponse(
+                    id=self.task_id,
+                    status_code=None,
+                    is_error=True,
+                    error_message=f"Unexpected {type(e).__name__}: {str(e) or 'No message.'}",
+                    texts=[],
+                    embeddings=[],
+                )
+            )
+            self.handle_error()
+
+
+@dataclass
+class EmbeddingResponse:
+    id: int
+    status_code: int | None
+    is_error: bool
+    error_message: Optional[str]
+    texts: list[str]
+    embeddings: list[list[float]]
+
+
+async def embed_parallel_async(
+    texts: list[str],
+    model: str = "rerank-english-v3.0",
+    max_attempts: int = 5,
+    max_requests_per_minute: int = 4_000,
+    max_concurrent_requests: int = 500,
+    request_timeout: int = 10,
+    batch_size: int = 16,
+    show_progress: bool = True,
+    **kwargs,
+):
+    """Processes embed requests in parallel, throttling to stay under rate limits."""
+    if batch_size > 96:
+        raise ValueError("Embeddings only support up to 96 texts per request.")
+    batches = [texts[i : i + batch_size] for i in range(0, len(texts), batch_size)]
+    pbar = tqdm(total=len(batches), desc="Embedding") if show_progress else None
+    ids = range(len(batches))
+    # constants
+    seconds_to_pause_after_rate_limit_error = 5
+    seconds_to_sleep_each_loop = 0.003  # so concurrent tasks can run
+
+    # initialize trackers
+    retry_queue = asyncio.Queue()
+    status_tracker = StatusTracker()
+    next_request = None  # variable to hold the next request to call
+
+    # initialize available capacity counts
+    # throttle over a 1 second window rather than minute,
+    # since some models limit RPS rather than RPM
+    available_request_capacity = max_requests_per_minute
+    last_update_time = time.time()
+    last_pbar_update_time = time.time()
+
+    # initialize flags
+    prompts_not_finished = True
+    prompts_iter = iter(zip(ids, batches))
+    results: list = []
+    session = aiohttp.ClientSession()
+
+    while True:
+        # get next request (if one is not already waiting for capacity)
+        if next_request is None:
+            if not retry_queue.empty():
+                next_request = retry_queue.get_nowait()
+                print(f"Retrying request {next_request.task_id}.")
+            elif prompts_not_finished:
+                try:
+                    # get new request
+                    batch_id, batch = next(prompts_iter)
+                    next_request = EmbeddingRequest(
+                        task_id=batch_id,
+                        model_name=model,
+                        texts=batch,
+                        attempts_left=max_attempts,
+                        status_tracker=status_tracker,
+                        retry_queue=retry_queue,
+                        request_timeout=request_timeout,
+                        pbar=pbar,
+                        **kwargs,
+                    )
+                    status_tracker.num_tasks_started += 1
+                    status_tracker.num_tasks_in_progress += 1
+                    results.append(next_request)
+
+                except StopIteration:
+                    prompts_not_finished = False
+                    print("API requests finished, only retries remain.")
+
+        # update available capacity
+        current_time = time.time()
+        seconds_since_update = current_time - last_update_time
+        available_request_capacity = min(
+            available_request_capacity
+            + max_requests_per_minute * seconds_since_update / 60.0,
+            max_requests_per_minute,
+        )
+        last_update_time = current_time
+
+        # update pbar status
+        if pbar:
+            if current_time - last_pbar_update_time > 1:
+                last_pbar_update_time = current_time
+                pbar.set_postfix(
+                    {
+                        "Req. Capacity": f"{available_request_capacity:.1f}",
+                        "Reqs in Progress": status_tracker.num_tasks_in_progress,
+                    }
+                )
+
+        # if enough capacity available, call API
+        if next_request:
+            if (
+                available_request_capacity >= 1
+                and status_tracker.num_tasks_in_progress < max_concurrent_requests
+            ):
+                # update counters
+                available_request_capacity -= 1
+                next_request.attempts_left -= 1
+
+                # call API
+                asyncio.create_task(next_request.call_api(session=session))
+                next_request = None  # reset next_request to empty
+
+        # if all tasks are finished, break
+        if status_tracker.num_tasks_in_progress == 0:
+            break
+
+        # main loop sleeps briefly so concurrent tasks can run
+        await asyncio.sleep(seconds_to_sleep_each_loop)
+
+        # if a rate limit error was hit recently, pause to cool down
+        seconds_since_rate_limit_error = (
+            time.time() - status_tracker.time_of_last_rate_limit_error
+        )
+        if seconds_since_rate_limit_error < seconds_to_pause_after_rate_limit_error:
+            remaining_seconds_to_pause = (
+                seconds_to_pause_after_rate_limit_error - seconds_since_rate_limit_error
+            )
+            await asyncio.sleep(remaining_seconds_to_pause)
+            # ^e.g., if pause is 15 seconds and final limit was hit 5 seconds ago
+            print(
+                f"Pausing to cool down until {time.ctime(status_tracker.time_of_last_rate_limit_error + seconds_to_pause_after_rate_limit_error)}"
+            )
+
+    # after finishing, log final status
+    if status_tracker.num_tasks_failed > 0:
+        print(
+            f"{status_tracker.num_tasks_failed} / {status_tracker.num_tasks_started} requests failed."
+        )
+    if status_tracker.num_rate_limit_errors > 0:
+        print(
+            f"{status_tracker.num_rate_limit_errors} rate limit errors received. Consider running at a lower rate."
+        )
+
+    print(
+        f"After processing, got {len(results)} results for {len(ids)} inputs. Removing duplicates."
+    )
+
+    # deduplicate results by id
+    deduplicated = {}
+    for request in results:
+        if request.task_id not in deduplicated:
+            deduplicated[request.task_id] = request.result[-1]
+        else:
+            current_response: EmbeddingResponse = deduplicated[request.task_id]
+            # only replace if the current request has no top_k_indexes and the new one does
+            if request.result[-1].embeddings and not current_response.embeddings:
+                deduplicated[request.task_id] = request.result[-1]
+
+    output = list(deduplicated.values())
+    # sort by id
+    output.sort(key=lambda x: x.id)
+    print(f"Returning {len(output)} unique results.")
+    await session.close()
+    return output
+
+
+def stack_results(
+    results: list[EmbeddingResponse], return_numpy: bool = True
+) -> list[list[float]] | np.ndarray:
+    if not all(response.status_code == 200 for response in results):
+        raise ValueError("Some responses were not successful; cannot coalesce results.")
+    stacked = np.concatenate([response.embeddings for response in results], axis=0)
+    return stacked.tolist() if not return_numpy else stacked  # type: ignore
+
+
+def submit_batch_request():
+    pass

lm_deluge/errors.py

@@ -0,0 +1,8 @@
+def raise_if_modal_exception(e: Exception):
+    try:
+        import modal  # type: ignore
+
+        if isinstance(e, modal.exception.InputCancellation):
+            raise e
+    except ImportError:
+        pass

lm_deluge/gemini_limits.py

@@ -0,0 +1,65 @@
+gemini_flash_limits = {
+    "asia-east1": 2000,
+    "asia-east2": 200,
+    "asia-northeast1": 200,
+    "asia-northeast3": 200,
+    "asia-south1": 200,
+    "asia-southeast1": 3_000,
+    "australia-southeast1": 200,
+    "europe-central2": 200,
+    "europe-north1": 200,
+    "europe-southwest1": 200,
+    "europe-west1": 10_000,
+    "europe-west2": 200,
+    "europe-west3": 200,
+    "europe-west4": 200,
+    "europe-west6": 200,
+    "europe-west8": 200,
+    "europe-west9": 200,
+    # 'me-central1': 200,
+    "me-central2": 200,
+    "me-west1": 200,
+    "northamerica-northeast1": 200,
+    "southamerica-east1": 200,
+    "us-central1": 5_000,
+    "us-east1": 3_000,
+    "us-east4": 200,
+    # 'us-east5': 200,
+    "us-south1": 3_000,
+    "us-west1": 5_000,
+    "us-west4": 200,
+}
+
+# total: 7_520
+gemini_1_5_pro_limits = {
+    "asia-east1": 500,
+    "asia-east2": 500,
+    "asia-northeast1": 500,
+    # "asia-northeast2": 500,
+    "asia-northeast3": 500,
+    "asia-south1": 500,
+    "asia-southeast1": 500,
+    "australia-southeast1": 60,
+    "europe-central2": 500,
+    "europe-north1": 60,
+    "europe-southwest1": 60,
+    "europe-west1": 500,
+    "europe-west2": 60,
+    "europe-west3": 60,
+    "europe-west4": 60,
+    "europe-west6": 60,
+    "europe-west8": 60,
+    "europe-west9": 60,
+    "me-central1": 60,
+    "me-central2": 60,
+    "me-west1": 60,
+    "northamerica-northeast1": 60,
+    "southamerica-east1": 500,
+    "us-central1": 500,
+    "us-east1": 500,
+    "us-east4": 60,
+    # "us-east5": 60,
+    "us-south1": 60,
+    "us-west1": 500,
+    "us-west4": 60,
+}

lm_deluge/image.py

@@ -0,0 +1,200 @@
+import os
+from contextlib import contextmanager
+from functools import cached_property
+import io
+import requests
+from PIL import Image as PILImage  # type: ignore
+import base64
+import mimetypes
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Literal
+
+
+@dataclass(slots=True)
+class Image:
+    # raw bytes, pathlike, http url, or base64 data url
+    data: bytes | io.BytesIO | Path | str
+    media_type: str | None = None  # inferred if None
+    detail: Literal["low", "high", "auto"] = "auto"
+    type: str = field(init=False, default="image")
+
+    @classmethod
+    def from_pdf(
+        cls,
+        pdf_path: str,
+        dpi: int = 200,
+        target_size: int = 1024,
+        first_page: int | None = None,
+        last_page: int | None = None,
+    ) -> list["Image"]:
+        try:
+            from pdf2image import convert_from_path  # type: ignore
+        except ImportError:
+            raise RuntimeError("pdf2image is required for PDF conversion.")
+
+        # Convert the first page of the PDF to an image
+        pages = convert_from_path(
+            pdf_path,
+            dpi=dpi,
+            first_page=first_page or 1,
+            last_page=last_page,  # type: ignore
+        )
+        images = []
+        for page in pages:
+            buffer = io.BytesIO()
+            page.save(buffer, format="JPEG")
+            image = cls(buffer.getvalue(), media_type="image/jpeg")
+            image.resize(target_size)
+            images.append(image)
+        return images
+
+    # helpers -----------------------------------------------------------------
+    def _bytes(self) -> bytes:
+        if isinstance(self.data, bytes):
+            return self.data
+        elif isinstance(self.data, io.BytesIO):
+            return self.data.getvalue()
+        elif isinstance(self.data, str) and self.data.startswith("http"):
+            res = requests.get(self.data)
+            res.raise_for_status()
+            return res.content
+        elif isinstance(self.data, str) and os.path.exists(self.data):
+            with open(self.data, "rb") as f:
+                return f.read()
+        elif isinstance(self.data, Path) and self.data.exists():
+            return Path(self.data).read_bytes()
+        elif isinstance(self.data, str) and self.data.startswith("data:"):
+            header, encoded = self.data.split(",", 1)
+            return base64.b64decode(encoded)
+        else:
+            raise ValueError("unreadable image format")
+
+    def _mime(self) -> str:
+        if self.media_type:
+            return self.media_type
+        if isinstance(self.data, (Path, str)):
+            guess = mimetypes.guess_type(str(self.data))[0]
+            if guess:
+                return guess
+        return "image/png"
+
+    def _base64(self, include_header: bool = True) -> str:
+        encoded = base64.b64encode(self._bytes()).decode("utf-8")
+        if not include_header:
+            return encoded
+        return f"data:{self._mime()};base64,{encoded}"
+
+    @contextmanager
+    def _image(self):
+        img = None
+        try:
+            img = PILImage.open(io.BytesIO(self._bytes()))
+            yield img
+        finally:
+            if img:
+                img.close()
+
+    @cached_property
+    def size(self) -> tuple[int, int]:
+        with self._image() as img:
+            return img.size
+
+    @cached_property
+    def num_pixels(self) -> int:
+        return self.size[0] * self.size[1]
+
+    def _resize(self, size: tuple[int, int]) -> bytes:
+        buffer = io.BytesIO()
+        new_width, new_height = size
+        with self._image() as img:
+            # Resize with Lanczos antialiasing
+            img.resize((new_width, new_height), PILImage.Resampling.LANCZOS).save(
+                buffer, format=self._mime().split("/")[-1].upper()
+            )
+
+        return buffer.getvalue()
+
+    def _resize_longer(
+        self, *, size: int | None = None, max_size: int | None = None
+    ) -> bytes:
+        if not max_size and not size:
+            raise ValueError("Either size or max_size must be provided")
+        width, height = self.size
+        if width > height:
+            new_width = size if size is not None else min(max_size, width)  # type: ignore
+            new_height = int(new_width / width * height)
+        else:
+            new_height = size if size is not None else min(max_size, height)  # type: ignore
+            new_width = int(new_height / height * width)
+        return self._resize((new_width, new_height))
+
+    def _resize_shorter(
+        self, *, size: int | None = None, max_size: int | None = None
+    ) -> bytes:
+        if not max_size and not size:
+            raise ValueError("Either size or max_size must be provided")
+        width, height = self.size
+        if width <= height:
+            new_width = size if size is not None else min(max_size, width)  # type: ignore
+            new_height = int(new_width / width * height)
+        else:
+            new_height = size if size is not None else min(max_size, height)  # type: ignore
+            new_width = int(new_height / height * width)
+        return self._resize((new_width, new_height))
+
+    @cached_property
+    def fingerprint(self) -> str:
+        # return base64 of a very small version of the image
+        small_image = self._resize_longer(max_size=48)  # longer side = 48px
+        return base64.b64encode(small_image).decode("utf-8")
+
+    def resize(self, max_size: int) -> None:
+        """
+        Resize the image and save to the data value.
+        """
+        self.data = self._resize_longer(max_size=max_size)
+
+    # ── provider-specific emission ────────────────────────────────────────────
+    def oa_chat(self) -> dict:
+        # if max(self.size) > 1_568:
+        #     self.resize_longer_side(1_568)
+        return {
+            "type": "image_url",
+            "image_url": {
+                "url": self._base64(),
+                "detail": self.detail,
+            },
+        }
+
+    def oa_resp(self) -> dict:
+        # if max(self.size) > 1_568:
+        #     self.resize_longer_side(1_568)
+        return {"type": "input_image", "image_url": self._base64()}
+
+    def anthropic(self) -> dict:
+        # n_pixels = self.num_pixels
+        # if n_pixels > 1_200_000:
+        #     resize_factor = (1_200_000 / n_pixels) ** 0.5
+        #     new_size = (
+        #         int(self.size[0] * resize_factor),
+        #         int(self.size[1] * resize_factor),
+        #     )
+        #     self.resize(new_size)
+        b64 = base64.b64encode(self._bytes()).decode()
+        return {
+            "type": "image",
+            "source": {
+                "type": "base64",
+                "media_type": self._mime(),
+                "data": b64,
+            },
+        }
+
+    def gemini(self) -> dict:
+        return {
+            "inlineData": {
+                "mimeType": self._mime(),
+                "data": self._base64(include_header=False),
+            }
+        }

lm_deluge/llm_tools/__init__.py

@@ -0,0 +1,11 @@
+from .extract import extract, extract_async
+from .translate import translate, translate_async
+from .score import score_llm
+
+__all__ = [
+    "extract",
+    "extract_async",
+    "translate",
+    "translate_async",
+    "score_llm",
+]