lm-deluge 0.0.4__py3-none-any.whl → 0.0.6__py3-none-any.whl

lm_deluge/api_requests/{cohere.py → mistral.py}

@@ -1,20 +1,19 @@
-# https://docs.cohere.com/reference/chat
-# https://cohere.com/pricing
 import asyncio
+import warnings
 from aiohttp import ClientResponse
 import json
 import os
-from tqdm import tqdm
+from tqdm.auto import tqdm
 from typing import Callable
-from lm_deluge.prompt import Conversation
-from .base import APIRequestBase, APIResponse
 
+from .base import APIRequestBase, APIResponse
+from ..prompt import Conversation
 from ..tracker import StatusTracker
 from ..sampling_params import SamplingParams
 from ..models import APIModel
 
 
-class CohereRequest(APIRequestBase):
+class MistralRequest(APIRequestBase):
     def __init__(
         self,
         task_id: int,
@@ -24,10 +23,12 @@ class CohereRequest(APIRequestBase):
         prompt: Conversation,
         attempts_left: int,
         status_tracker: StatusTracker,
-        results_arr: list,
         retry_queue: asyncio.Queue,
+        results_arr: list,
         request_timeout: int = 30,
         sampling_params: SamplingParams = SamplingParams(),
+        logprobs: bool = False,
+        top_logprobs: int | None = None,
         pbar: tqdm | None = None,
         callback: Callable | None = None,
         debug: bool = False,
@@ -44,36 +45,36 @@ class CohereRequest(APIRequestBase):
             results_arr=results_arr,
             request_timeout=request_timeout,
             sampling_params=sampling_params,
+            logprobs=logprobs,
+            top_logprobs=top_logprobs,
             pbar=pbar,
             callback=callback,
             debug=debug,
             all_model_names=all_model_names,
             all_sampling_params=all_sampling_params,
         )
-        self.system_message = None
-        self.last_user_message = None
-
         self.model = APIModel.from_registry(model_name)
-        self.url = f"{self.model.api_base}/chat"
-        self.system_message, chat_history, last_user_message = prompt.to_cohere()
-
+        self.url = f"{self.model.api_base}/chat/completions"
         self.request_header = {
-            "Authorization": f"bearer {os.getenv(self.model.api_key_env_var)}",
-            "content-type": "application/json",
-            "accept": "application/json",
+            "Authorization": f"Bearer {os.getenv(self.model.api_key_env_var)}"
         }
-
         self.request_json = {
             "model": self.model.name,
-            "chat_history": chat_history,
-            "message": last_user_message,
+            "messages": prompt.to_mistral(),
             "temperature": sampling_params.temperature,
             "top_p": sampling_params.top_p,
             "max_tokens": sampling_params.max_new_tokens,
         }
-
-        if self.system_message:
-            self.request_json["preamble"] = self.system_message
+        if sampling_params.reasoning_effort:
+            warnings.warn(
+                f"Ignoring reasoning_effort param for non-reasoning model: {model_name}"
+            )
+        if logprobs:
+            warnings.warn(
+                f"Ignoring logprobs param for non-logprobs model: {model_name}"
+            )
+        if sampling_params.json_mode and self.model.supports_json:
+            self.request_json["response_format"] = {"type": "json_object"}
 
     async def handle_response(self, http_response: ClientResponse) -> APIResponse:
         is_error = False
@@ -81,41 +82,41 @@ class CohereRequest(APIRequestBase):
         completion = None
         input_tokens = None
         output_tokens = None
+        logprobs = None
         status_code = http_response.status
         mimetype = http_response.headers.get("Content-Type", None)
+        data = None
         if status_code >= 200 and status_code < 300:
             try:
                 data = await http_response.json()
             except Exception:
-                data = None
                 is_error = True
                 error_message = (
                     f"Error calling .json() on response w/ status {status_code}"
                 )
-            if not is_error and isinstance(data, dict):
+            if not is_error:
+                assert data is not None, "data is None"
                 try:
-                    completion = data["text"]
-                    input_tokens = data["meta"]["billed_units"]["input_tokens"]
-                    output_tokens = data["meta"]["billed_units"]["input_tokens"]
+                    completion = data["choices"][0]["message"]["content"]
+                    input_tokens = data["usage"]["prompt_tokens"]
+                    output_tokens = data["usage"]["completion_tokens"]
+                    if self.logprobs and "logprobs" in data["choices"][0]:
+                        logprobs = data["choices"][0]["logprobs"]["content"]
                 except Exception:
                     is_error = True
-                    error_message = f"Error getting 'text' or 'meta' from {self.model.name} response."
-        elif mimetype is not None and "json" in mimetype.lower():
+                    error_message = f"Error getting 'choices' and 'usage' from {self.model.name} response."
+        elif mimetype and "json" in mimetype.lower():
             is_error = True  # expected status is 200, otherwise it's an error
             data = await http_response.json()
             error_message = json.dumps(data)
-
         else:
             is_error = True
             text = await http_response.text()
             error_message = text
 
-        # handle special kinds of errors. TODO: make sure these are correct for anthropic
+        # handle special kinds of errors
         if is_error and error_message is not None:
-            if (
-                "rate limit" in error_message.lower()
-                or "overloaded" in error_message.lower()
-            ):
+            if "rate limit" in error_message.lower() or status_code == 429:
                 error_message += " (Rate limit error, triggering cooldown.)"
                 self.status_tracker.rate_limit_exceeded()
             if "context length" in error_message:
@@ -128,6 +129,7 @@ class CohereRequest(APIRequestBase):
             is_error=is_error,
             error_message=error_message,
             prompt=self.prompt,
+            logprobs=logprobs,
             completion=completion,
             model_internal=self.model_name,
             sampling_params=self.sampling_params,

lm_deluge/api_requests/openai.py

@@ -58,13 +58,18 @@ class OpenAIRequest(APIRequestBase):
         self.request_header = {
             "Authorization": f"Bearer {os.getenv(self.model.api_key_env_var)}"
         }
+
         self.request_json = {
             "model": self.model.name,
             "messages": prompt.to_openai(),
             "temperature": sampling_params.temperature,
             "top_p": sampling_params.top_p,
-            "max_completion_tokens": sampling_params.max_new_tokens,
         }
+        # set max_tokens or max_completion_tokens dep. on provider
+        if "cohere" in self.model.api_base:
+            self.request_json["max_tokens"] = sampling_params.max_new_tokens
+        elif "openai" in self.model.api_base:
+            self.request_json["max_completion_tokens"] = sampling_params.max_new_tokens
         if self.model.reasoning_model:
             self.request_json["temperature"] = 1.0
             self.request_json["top_p"] = 1.0
@@ -84,6 +89,7 @@ class OpenAIRequest(APIRequestBase):
     async def handle_response(self, http_response: ClientResponse) -> APIResponse:
         is_error = False
         error_message = None
+        thinking = None
         completion = None
         input_tokens = None
         output_tokens = None
@@ -103,6 +109,8 @@ class OpenAIRequest(APIRequestBase):
                 assert data is not None, "data is None"
                 try:
                     completion = data["choices"][0]["message"]["content"]
+                    if "reasoning_content" in data["choices"][0]["message"]:
+                        thinking = data["choices"][0]["message"]["reasoning_content"]
                     input_tokens = data["usage"]["prompt_tokens"]
                     output_tokens = data["usage"]["completion_tokens"]
                     if self.logprobs and "logprobs" in data["choices"][0]:
@@ -135,6 +143,7 @@ class OpenAIRequest(APIRequestBase):
             error_message=error_message,
             prompt=self.prompt,
             logprobs=logprobs,
+            thinking=thinking,
             completion=completion,
             model_internal=self.model_name,
             sampling_params=self.sampling_params,

lm_deluge/client.py

@@ -85,6 +85,7 @@ class LLMClient:
     def __init__(
         self,
         model_names: list[str],
+        *,
         max_requests_per_minute: int,
         max_tokens_per_minute: int,
         max_concurrent_requests: int,
@@ -345,6 +346,7 @@ class LLMClient:
 
         # add cache hits back in
         for id, res in zip(cache_hit_ids, cache_hit_results):
+            res.cache_hit = True
             results[id] = res
 
         if return_completions_only:

lm_deluge/image.py

@@ -191,6 +191,12 @@ class Image:
             },
         }
 
+    def mistral(self) -> dict:
+        return {
+            "type": "image_url",
+            "image_url": self._base64(),
+        }
+
     def gemini(self) -> dict:
         return {
             "inlineData": {