speedy-utils 1.1.26__py3-none-any.whl → 1.1.28__py3-none-any.whl

llm_utils/__init__.py

@@ -1,8 +1,19 @@
+from llm_utils.lm import (
+    LLM,
+    AsyncLLMTask,
+    AsyncLM,
+    Input,
+    InputField,
+    LLMSignature,
+    Output,
+    OutputField,
+    Signature,
+)
+from llm_utils.lm.base_prompt_builder import BasePromptBuilder
+from llm_utils.lm.lm_base import get_model_name
 from llm_utils.lm.openai_memoize import MOpenAI
-from llm_utils.lm import LLM, AsyncLM, AsyncLLMTask, LLMSignature, Signature, InputField, OutputField, Input, Output
 from llm_utils.vector_cache import VectorCache
-from llm_utils.lm.lm_base import get_model_name
-from llm_utils.lm.base_prompt_builder import BasePromptBuilder
+
 
 LLM_TASK = LLM
 
@@ -24,13 +35,14 @@ from llm_utils.chat_format import (
     display_conversations,
     format_msgs,
     get_conversation_one_turn,
-    show_chat_v2,
     show_chat,
+    show_chat_v2,
     show_string_diff,
     transform_messages,
     transform_messages_to_chatml,
 )
 
+
 __all__ = [
     "transform_messages",
     "transform_messages_to_chatml",

llm_utils/chat_format/__init__.py

@@ -1,17 +1,17 @@
-from .transform import (
-    identify_format,
-    _transform_sharegpt_to_chatml,
-    transform_messages,
-    transform_messages_to_chatml,
-)
 from .display import (
-    show_chat,
+    display_chat_messages_as_html,
+    display_conversations,
     get_conversation_one_turn,
     highlight_diff_chars,
-    show_string_diff,
-    display_conversations,
-    display_chat_messages_as_html,
+    show_chat,
     show_chat_v2,
+    show_string_diff,
+)
+from .transform import (
+    _transform_sharegpt_to_chatml,
+    identify_format,
+    transform_messages,
+    transform_messages_to_chatml,
 )
 from .utils import (
     build_chatml_input,

llm_utils/chat_format/display.py

@@ -77,7 +77,7 @@ def show_chat(
     theme: str = "default",
     as_markdown: bool = False,
     as_json: bool = False,
-) -> Optional[str]:
+) -> str | None:
     """
     Display chat messages as HTML.
 
@@ -168,7 +168,10 @@ def show_chat(
             content = content.replace("\t", "    ")
             content = content.replace("  ", "  ")
             content = (
-                content.replace("<br>", "TEMP_BR").replace("<", "&lt;").replace(">", "&gt;").replace("TEMP_BR", "<br>")
+                content.replace("<br>", "TEMP_BR")
+                .replace("<", "&lt;")
+                .replace(">", "&gt;")
+                .replace("TEMP_BR", "<br>")
             )
         if role in color_scheme:
             background_color = color_scheme[role]["background"]
@@ -239,15 +242,15 @@ def show_chat(
             f.write(html)
     if return_html:
         return html
-    else:
-        display(HTML(html))
+    display(HTML(html))
+    return None
 
 
 def get_conversation_one_turn(
-    system_msg: Optional[str] = None,
-    user_msg: Optional[str] = None,
-    assistant_msg: Optional[str] = None,
-    assistant_prefix: Optional[str] = None,
+    system_msg: str | None = None,
+    user_msg: str | None = None,
+    assistant_msg: str | None = None,
+    assistant_prefix: str | None = None,
     return_format: str = "chatml",
 ) -> Any:
     """
@@ -261,7 +264,9 @@ def get_conversation_one_turn(
     if assistant_msg is not None:
         messages.append({"role": "assistant", "content": assistant_msg})
     if assistant_prefix is not None:
-        assert return_format != "chatml", 'Change return_format to "text" if you want to use assistant_prefix'
+        assert (
+            return_format != "chatml"
+        ), 'Change return_format to "text" if you want to use assistant_prefix'
         assert messages[-1]["role"] == "user"
         from .transform import transform_messages
 
@@ -270,9 +275,8 @@ def get_conversation_one_turn(
             msg = str(msg)
         msg += assistant_prefix
         return msg
-    else:
-        assert return_format in ["chatml"]
-        return messages
+    assert return_format in ["chatml"]
+    return messages
 
 
 def highlight_diff_chars(text1: str, text2: str) -> str:
@@ -286,13 +290,21 @@ def highlight_diff_chars(text1: str, text2: str) -> str:
             html.append(text1[i1:i2])
         elif tag == "replace":
             if i1 != i2:
-                html.append(f'<span style="background-color:#ffd6d6; color:#b20000;">{text1[i1:i2]}</span>')
+                html.append(
+                    f'<span style="background-color:#ffd6d6; color:#b20000;">{text1[i1:i2]}</span>'
+                )
             if j1 != j2:
-                html.append(f'<span style="background-color:#d6ffd6; color:#006600;">{text2[j1:j2]}</span>')
+                html.append(
+                    f'<span style="background-color:#d6ffd6; color:#006600;">{text2[j1:j2]}</span>'
+                )
         elif tag == "delete":
-            html.append(f'<span style="background-color:#ffd6d6; color:#b20000;">{text1[i1:i2]}</span>')
+            html.append(
+                f'<span style="background-color:#ffd6d6; color:#b20000;">{text1[i1:i2]}</span>'
+            )
         elif tag == "insert":
-            html.append(f'<span style="background-color:#d6ffd6; color:#006600;">{text2[j1:j2]}</span>')
+            html.append(
+                f'<span style="background-color:#d6ffd6; color:#006600;">{text2[j1:j2]}</span>'
+            )
     return "".join(html)
 
 
@@ -321,7 +333,7 @@ def show_chat_v2(messages: list[dict[str, str]]):
 
     if is_notebook:
         # Use HTML display in notebook
-        from IPython.display import display, HTML
+        from IPython.display import HTML, display
 
         role_colors = {
             "system": "red",
@@ -353,9 +365,7 @@ def show_chat_v2(messages: list[dict[str, str]]):
             html += f"<div style='color:{color}'><strong>{label}</strong><br>{content}</div>"
             # Add separator except after last message
             if i < len(messages) - 1:
-                html += (
-                    "<div style='color:#888; margin:0.5em 0;'>───────────────────────────────────────────────────</div>"
-                )
+                html += "<div style='color:#888; margin:0.5em 0;'>───────────────────────────────────────────────────</div>"
         html += "</div>"
 
         display(HTML(html))
@@ -385,7 +395,9 @@ def show_chat_v2(messages: list[dict[str, str]]):
             print(f"{color}{content}{reset}")
             # Add separator except after last message
             if i < len(messages) - 1:
-                print(f"{separator_color}─────────────────────────────────────────────────────────{reset}")
+                print(
+                    f"{separator_color}─────────────────────────────────────────────────────────{reset}"
+                )
 
 
 def display_conversations(data1: Any, data2: Any, theme: str = "light") -> None:

llm_utils/chat_format/transform.py

@@ -1,13 +1,13 @@
 from __future__ import annotations
+
 from copy import deepcopy
 
 
 def identify_format(item):
     if isinstance(item, list) and "role" in item[0]:
         return "chatml"
-    if isinstance(item, dict):
-        if "conversations" in item:
-            return "sharegpt"
+    if isinstance(item, dict) and "conversations" in item:
+        return "sharegpt"
     raise ValueError(
         f"The format of the item is not recognized. \n{type(item)=}, \n{item=}"
     )
@@ -16,9 +16,9 @@ def identify_format(item):
 def _transform_sharegpt_to_chatml(
     item, default_system_message="You are a helpful assistant.", print_msg=False
 ):
-    assert isinstance(item, dict), (
-        "The item is not in the correct format. Please check the format of the item."
-    )
+    assert isinstance(
+        item, dict
+    ), "The item is not in the correct format. Please check the format of the item."
 
     messages = []
     system_msg = item.get("system", "")
@@ -82,16 +82,16 @@ def transform_messages(
                         {"from": message["role"], "value": message["content"]}
                     )
             return ret
-        elif to == "chatml":
+        if to == "chatml":
             return _transform_sharegpt_to_chatml(item)
-        elif to == "text":
+        if to == "text":
             text = ""
             for turn in chatml_messages:
                 text += f"<|im_start|>{turn['role']}\n{turn['content']}<|im_end|>\n"
             if add_generation_prompt:
                 text += "<|im_start|>assistant\n"
             return text
-        elif to == "simulated_chat":
+        if to == "simulated_chat":
             text = "<role> Given the simulated chat, you are the assistant. Lets continue the conversation. \n\n"
             for turn in chatml_messages:
                 prefix = {
@@ -104,11 +104,9 @@ def transform_messages(
             if add_generation_prompt:
                 text += "AI: [continue the conversation here]"
             return text
-        else:
-            raise ValueError(f"{to} is not supported.")
+        raise ValueError(f"{to} is not supported.")
 
-    else:
-        return item
+    return item
 
 
 def transform_messages_to_chatml(input_data, input_format="auto"):
@@ -116,16 +114,16 @@ def transform_messages_to_chatml(input_data, input_format="auto"):
         input_data = deepcopy(input_data)
         if isinstance(input_data, list):
             input_format = "chatlm"
-            assert input_data[0].get("role") is not None, (
-                "The input format is not recognized. Please specify the input format."
-            )
+            assert (
+                input_data[0].get("role") is not None
+            ), "The input format is not recognized. Please specify the input format."
         elif isinstance(input_data, dict):
             input_data = _transform_sharegpt_to_chatml(input_data)
             input_format = "sharegpt"
         elif isinstance(input_data, str):
-            assert "<|im_end|>" in input_data, (
-                "The input format is not recognized. Please specify the input format."
-            )
+            assert (
+                "<|im_end|>" in input_data
+            ), "The input format is not recognized. Please specify the input format."
             input_format = "chatlm"
             parts = input_data.split("<|im_end|>")
             input_data = []

llm_utils/chat_format/utils.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
-from typing import List, Dict, Callable
 
+from collections.abc import Callable
+from typing import Dict, List
 
-def build_chatml_input(template: str, params: List[str]) -> Callable:
-    def formator(**kwargs) -> List[List[Dict[str, str]]]:
-        system_msg = kwargs.get("system_msg", None)
+
+def build_chatml_input(template: str, params: list[str]) -> Callable:
+    def formator(**kwargs) -> list[list[dict[str, str]]]:
+        system_msg = kwargs.get("system_msg")
         kwargs.pop("system_msg", None)
         for param in params:
             if param not in kwargs:

llm_utils/group_messages.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 import random
-from typing import Sequence, Optional
+from collections.abc import Sequence
+from typing import Optional, cast
 
 import numpy as np
 import pandas as pd
@@ -16,7 +17,7 @@ def split_indices_by_length(
     random_seed: int,
     verbose: bool,
     shuffle: bool,
-    mean_length: Optional[int] = None,
+    mean_length: int | None = None,
 ) -> list[list[int]]:
     """
     Split indices into batches so that the sum of lengths in each batch does not exceed max_batch_length.
@@ -55,19 +56,19 @@ def split_indices_by_length(
         desc = pd.Series(batch_lengths).describe()
 
         table = [
-            ["New avg item len", desc["mean"]],
-            ["Number groups", len(batches)],
-            ["Max length", max_batch_length],
+            ['New avg item len', desc['mean']],
+            ['Number groups', len(batches)],
+            ['Max length', max_batch_length],
         ]
 
-        print(tabulate(table, headers=["Metric", "Value"], tablefmt="pretty"))
+        print(tabulate(table, headers=['Metric', 'Value'], tablefmt='pretty'))
 
     return batches
 
 
 def group_messages_by_len(
     messages: Sequence[dict],
-    model_name: str = "Qwen/Qwen2.5-7B-Instruct",
+    model_name: str = 'Qwen/Qwen2.5-7B-Instruct',
     batch_size: int = 4,
     mean_length: int = 512,
 ) -> list[dict]:
@@ -75,17 +76,19 @@ def group_messages_by_len(
     Groups messages into batches based on token length and concatenates them.
     """
     if messages is None:
-        raise ValueError("messages parameter cannot be None")
+        raise ValueError('messages parameter cannot be None')
     from transformers.models.auto.tokenization_auto import AutoTokenizer  # type: ignore
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
     def create_batches(messages: Sequence[dict]) -> list[dict]:
         def get_token_length(message: dict) -> int:
-            ids = tokenizer.apply_chat_template(message["messages"][1:], tokenize=True)
+            ids = tokenizer.apply_chat_template(message['messages'][1:], tokenize=True)
             return len(ids)
 
-        lengths: list[int] = multi_thread(get_token_length, messages, workers=64)
+        lengths: list[int] = cast(
+            list[int], multi_thread(get_token_length, messages, workers=64)
+        )
         list_ids: list[list[int]] = split_indices_by_length(
             lengths,
             batch_size,
@@ -101,12 +104,12 @@ def group_messages_by_len(
             turns: list[dict] = []
             for conv in conversations:
                 turns.extend(conv[1:])
-            return {"messages": [system_message] + turns}
+            return {'messages': [system_message] + turns}
 
         for batch_ids in list_ids:
             if not batch_ids:
                 continue
-            conversations = [messages[i]["messages"] for i in batch_ids]
+            conversations = [messages[i]['messages'] for i in batch_ids]
             concatenated_messages.append(concatenate_messages(conversations))
         return concatenated_messages
 
@@ -115,6 +118,6 @@ def group_messages_by_len(
 
 
 __all__ = [
-    "split_indices_by_length",
-    "group_messages_by_len",
+    'split_indices_by_length',
+    'group_messages_by_len',
 ]

llm_utils/lm/__init__.py

@@ -1,16 +1,17 @@
-from .async_lm.async_lm import AsyncLM
 from .async_lm.async_llm_task import AsyncLLMTask
-from .lm_base import LMBase, get_model_name
-from .llm import LLM
+from .async_lm.async_lm import AsyncLM
 from .base_prompt_builder import BasePromptBuilder
+from .llm import LLM
 from .llm_signature import LLMSignature
-from .signature import Signature, InputField, OutputField, Input, Output
+from .lm_base import LMBase, get_model_name
 from .mixins import (
+    ModelUtilsMixin,
     TemperatureRangeMixin,
     TwoStepPydanticMixin,
     VLLMMixin,
-    ModelUtilsMixin,
 )
+from .signature import Input, InputField, Output, OutputField, Signature
+
 
 __all__ = [
     "LMBase",

llm_utils/lm/async_lm/__init__.py

@@ -1,6 +1,7 @@
 from .async_llm_task import AsyncLLMTask
 from .async_lm import AsyncLM
 
+
 __all__ = [
     "AsyncLM",
     "AsyncLLMTask",

llm_utils/lm/async_lm/_utils.py

@@ -15,12 +15,13 @@ from openai.types.chat import (
 from pydantic import BaseModel
 from typing_extensions import TypedDict
 
+
 # --------------------------------------------------------------------------- #
 # type helpers
 # --------------------------------------------------------------------------- #
 TModel = TypeVar("TModel", bound=BaseModel)
-Messages = List[ChatCompletionMessageParam]
-LegacyMsgs = List[Dict[str, str]]
+Messages = list[ChatCompletionMessageParam]
+LegacyMsgs = list[dict[str, str]]
 RawMsgs = Union[Messages, LegacyMsgs]
 
 # --------------------------------------------------------------------------- #
@@ -55,10 +56,10 @@ OutputModelType = TypeVar("OutputModelType", bound=BaseModel)
 
 
 class ParsedOutput(TypedDict, Generic[OutputModelType]):
-    messages: List
+    messages: list
     completion: Any
     parsed: OutputModelType
-    model_kwargs: Dict[str, Any]
+    model_kwargs: dict[str, Any]
 
 
 # --------------------------------------------------------------------------- #
@@ -83,7 +84,7 @@ async def inspect_word_probs_async(lm, tokenizer, messages):
     async def compute_word_log_probs(
         tokenizer: Any,
         lm_client: Any,
-    ) -> tuple[List[Dict[str, Any]], Any]:
+    ) -> tuple[list[dict[str, Any]], Any]:
         # Build a prompt that preserves literal newlines
         prompt = tokenizer.apply_chat_template(
             messages,
@@ -112,7 +113,7 @@ async def inspect_word_probs_async(lm, tokenizer, messages):
         }
 
         # Flatten tokens
-        tokens: List[Dict[str, Any]] = [
+        tokens: list[dict[str, Any]] = [
             {"id": int(tid), **tdata}
             for td in token_logprob_dicts
             for tid, tdata in td.items()
@@ -133,7 +134,7 @@ async def inspect_word_probs_async(lm, tokenizer, messages):
         split_prompt = prompt.replace("\n", " <NL> ")
         words = split_prompt.split()
 
-        word_log_probs: List[Dict[str, Any]] = []
+        word_log_probs: list[dict[str, Any]] = []
         token_idx = 0
 
         for word in words:
@@ -152,7 +153,7 @@ async def inspect_word_probs_async(lm, tokenizer, messages):
 
         return word_log_probs, token_logprob_dicts  # type: ignore
 
-    def render_by_logprob(word_log_probs: List[Dict[str, Any]]) -> str:
+    def render_by_logprob(word_log_probs: list[dict[str, Any]]) -> str:
         """
         Return an ANSI-colored string for word probabilities (red → green).
         """
@@ -161,7 +162,7 @@ async def inspect_word_probs_async(lm, tokenizer, messages):
 
         probs = [entry["probability"] for entry in word_log_probs]
         min_p, max_p = min(probs), max(probs)
-        parts: List[str] = []
+        parts: list[str] = []
 
         for entry in word_log_probs:
             word = entry["word"]