speedy-utils 1.0.4__py3-none-any.whl → 1.0.9__py3-none-any.whl

llm_utils/__init__.py

@@ -0,0 +1,31 @@
+from .chat_format import (
+    transform_messages,
+    transform_messages_to_chatml,
+    show_chat,
+    get_conversation_one_turn,
+    show_string_diff,
+    display_conversations,
+    build_chatml_input,
+    format_msgs,
+    display_chat_messages_as_html,
+)
+from .lm import LM
+from .group_messages import (
+    split_indices_by_length,
+    group_messages_by_len,
+)
+
+__all__ = [
+    "transform_messages",
+    "transform_messages_to_chatml",
+    "show_chat",
+    "get_conversation_one_turn",
+    "show_string_diff",
+    "display_conversations",
+    "build_chatml_input",
+    "format_msgs",
+    "split_indices_by_length",
+    "group_messages_by_len",
+    "LM",
+    "display_chat_messages_as_html",
+]

llm_utils/chat_format/__init__.py

@@ -0,0 +1,34 @@
+from .transform import (
+    identify_format,
+    _transform_sharegpt_to_chatml,
+    transform_messages,
+    transform_messages_to_chatml,
+)
+from .display import (
+    show_chat,
+    get_conversation_one_turn,
+    highlight_diff_chars,
+    show_string_diff,
+    display_conversations,
+    display_chat_messages_as_html,
+)
+from .utils import (
+    build_chatml_input,
+    format_msgs,
+)
+
+
+__all__ = [
+    "identify_format",
+    "_transform_sharegpt_to_chatml",
+    "transform_messages",
+    "transform_messages_to_chatml",
+    "show_chat",
+    "get_conversation_one_turn",
+    "highlight_diff_chars",
+    "build_chatml_input",
+    "format_msgs",
+    "show_string_diff",
+    "display_conversations",
+    "display_chat_messages_as_html",
+]

llm_utils/chat_format/display.py

@@ -0,0 +1,274 @@
+from __future__ import annotations
+from typing import List, Tuple, Sequence, Any, Dict, Optional
+from IPython.display import HTML, display
+from difflib import SequenceMatcher
+
+
+def show_chat(
+    msgs: Any,
+    return_html: bool = False,
+    file: str = "/tmp/conversation.html",
+    theme: str = "default",
+) -> Optional[str]:
+    """
+    Display chat messages as HTML.
+    """
+    if isinstance(msgs, dict) and "messages" in msgs:
+        msgs = msgs["messages"]
+    assert isinstance(msgs, list) and all(
+        isinstance(msg, dict) and "role" in msg and "content" in msg for msg in msgs
+    ), "The input format is not recognized. Please specify the input format."
+
+    themes: dict[str, dict[str, dict[str, str]]] = {
+        "default": {
+            "system": {"background": "#ffaaaa", "text": "#222222"},  # More red
+            "user": {"background": "#f8c57e", "text": "#222222"},  # More orange
+            "assistant": {"background": "#9dfebd", "text": "#222222"},  # More green
+            "function": {"background": "#eafde7", "text": "#222222"},
+            "tool": {"background": "#fde7fa", "text": "#222222"},
+            "default": {"background": "#ffffff", "text": "#222222"},
+        },
+        "light": {
+            "system": {"background": "#ff6666", "text": "#000000"},  # More red
+            "user": {"background": "#ffd580", "text": "#000000"},  # More orange
+            "assistant": {"background": "#80ffb3", "text": "#000000"},  # More green
+            "function": {"background": "#AFFFFF", "text": "#000000"},
+            "tool": {"background": "#FFAAFF", "text": "#000000"},
+            "default": {"background": "#FFFFFF", "text": "#000000"},
+        },
+        "dark": {
+            "system": {"background": "#b22222", "text": "#fffbe7"},  # More red
+            "user": {"background": "#ff8800", "text": "#18181b"},  # More orange
+            "assistant": {"background": "#22c55e", "text": "#e0ffe0"},  # More green
+            "function": {"background": "#134e4a", "text": "#e0fff7"},
+            "tool": {"background": "#701a75", "text": "#ffe0fa"},
+            "default": {"background": "#18181b", "text": "#f4f4f5"},
+        },
+    }
+
+    color_scheme = themes.get(theme, themes["default"])
+
+    conversation_html = ""
+    for i, message in enumerate(msgs):
+        role = message["role"]
+        content = message.get("content", "")
+        if not content:
+            content = ""
+        tool_calls = message.get("tool_calls")
+        if not content and tool_calls:
+            for tool_call in tool_calls:
+                tool_call = tool_call["function"]
+                name = tool_call["name"]
+                args = tool_call["arguments"]
+                content += f"Tool: {name}\nArguments: {args}"
+        content = content.replace("\n", "<br>")
+        content = content.replace("\t", "&nbsp;&nbsp;&nbsp;&nbsp;")
+        content = content.replace("  ", "&nbsp;&nbsp;")
+        content = (
+            content.replace("<br>", "TEMP_BR")
+            .replace("<", "&lt;")
+            .replace(">", "&gt;")
+            .replace("TEMP_BR", "<br>")
+        )
+        if role in color_scheme:
+            background_color = color_scheme[role]["background"]
+            text_color = color_scheme[role]["text"]
+        else:
+            background_color = color_scheme["default"]["background"]
+            text_color = color_scheme["default"]["text"]
+        if role == "system":
+            conversation_html += (
+                f'<div style="background-color: {background_color}; color: {text_color}; padding: 10px; margin-bottom: 10px;">'
+                f'<strong>System:</strong><br><pre id="system-{i}">{content}</pre></div>'
+            )
+        elif role == "user":
+            conversation_html += (
+                f'<div style="background-color: {background_color}; color: {text_color}; padding: 10px; margin-bottom: 10px;">'
+                f'<strong>User:</strong><br><pre id="user-{i}">{content}</pre></div>'
+            )
+        elif role == "assistant":
+            conversation_html += (
+                f'<div style="background-color: {background_color}; color: {text_color}; padding: 10px; margin-bottom: 10px;">'
+                f'<strong>Assistant:</strong><br><pre id="assistant-{i}">{content}</pre></div>'
+            )
+        elif role == "function":
+            conversation_html += (
+                f'<div style="background-color: {background_color}; color: {text_color}; padding: 10px; margin-bottom: 10px;">'
+                f'<strong>Function:</strong><br><pre id="function-{i}">{content}</pre></div>'
+            )
+        else:
+            conversation_html += (
+                f'<div style="background-color: {background_color}; color: {text_color}; padding: 10px; margin-bottom: 10px;">'
+                f'<strong>{role}:</strong><br><pre id="{role}-{i}">{content}</pre><br>'
+                f"<button onclick=\"copyContent('{role}-{i}')\">Copy</button></div>"
+            )
+    html: str = f"""
+    <html>
+    <head>
+        <style>
+            pre {{
+                white-space: pre-wrap;
+            }}
+        </style>
+    </head>
+    <body>
+        {conversation_html}
+        <script>
+            function copyContent(elementId) {{
+                var element = document.getElementById(elementId);
+                var text = element.innerText;
+                navigator.clipboard.writeText(text)
+                    .then(function() {{
+                        alert("Content copied to clipboard!");
+                    }})
+                    .catch(function(error) {{
+                        console.error("Error copying content: ", error);
+                    }});
+            }}
+        </script>
+    </body>
+    </html>
+    """
+    if file:
+        with open(file, "w") as f:
+            f.write(html)
+    if return_html:
+        return html
+    else:
+        display(HTML(html))
+
+
+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,
+    return_format: str = "chatml",
+) -> Any:
+    """
+    Build a one-turn conversation.
+    """
+    messages: list[dict[str, str]] = []
+    if system_msg is not None:
+        messages.append({"role": "system", "content": system_msg})
+    if user_msg is not None:
+        messages.append({"role": "user", "content": user_msg})
+    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 messages[-1]["role"] == "user"
+        from .transform import transform_messages
+
+        msg = transform_messages(messages, "chatml", "text", add_generation_prompt=True)
+        if not isinstance(msg, str):
+            msg = str(msg)
+        msg += assistant_prefix
+        return msg
+    else:
+        assert return_format in ["chatml"]
+        return messages
+
+
+def highlight_diff_chars(text1: str, text2: str) -> str:
+    """
+    Return a string with deletions in red and additions in green.
+    """
+    matcher = SequenceMatcher(None, text1, text2)
+    html: list[str] = []
+    for tag, i1, i2, j1, j2 in matcher.get_opcodes():
+        if tag == "equal":
+            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>'
+                )
+            if j1 != j2:
+                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>'
+            )
+        elif tag == "insert":
+            html.append(
+                f'<span style="background-color:#d6ffd6; color:#006600;">{text2[j1:j2]}</span>'
+            )
+    return "".join(html)
+
+
+def show_string_diff(old: str, new: str) -> None:
+    """
+    Display a one-line visual diff between two strings (old -> new).
+    """
+    html1 = highlight_diff_chars(old, new)
+    display(HTML(html1))
+
+
+def display_conversations(data1: Any, data2: Any, theme: str = "light") -> None:
+    """
+    Display two conversations side by side.
+    """
+    import warnings
+
+    warnings.warn(
+        "display_conversations will be deprecated in the next version.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    html1 = show_chat(data1, return_html=True, theme=theme)
+    html2 = show_chat(data2, return_html=True, theme=theme)
+    html = f"""
+    <html>
+    <head>
+        <style>
+            table {{
+                width: 100%;
+                border-collapse: collapse;
+            }}
+            td {{
+                width: 50%;
+                vertical-align: top;
+                padding: 10px;
+            }}
+        </style>
+    </head>
+    <body>
+        <table>
+            <tr>
+                <td>{html1}</td>
+                <td>{html2}</td>
+            </tr>
+        </table>
+    </body>
+    </html>
+    """
+    display(HTML(html))
+
+
+def display_chat_messages_as_html(*args, **kwargs):
+    """
+    Use as show_chat and warn about the deprecated function.
+    """
+    import warnings
+
+    warnings.warn(
+        "display_chat_messages_as_html is deprecated, use show_chat instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return show_chat(*args, **kwargs)
+
+
+__all__ = [
+    "show_chat",
+    "get_conversation_one_turn",
+    "highlight_diff_chars",
+    "show_string_diff",
+    "display_conversations",
+    "display_chat_messages_as_html",
+]

llm_utils/chat_format/transform.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+from copy import deepcopy
+from typing import Callable, Dict, List, Sequence
+
+
+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"
+    raise ValueError(
+        f"The format of the item is not recognized. \n{type(item)=}, \n{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."
+
+    messages = []
+    system_msg = item.get("system", "")
+    if system_msg:
+        messages.append({"role": "system", "content": system_msg})
+    elif default_system_message:
+        messages.append({"role": "system", "content": default_system_message})
+    conversations = item.get("conversations", [])
+    if hasattr(conversations, "tolist"):
+        conversations = conversations.tolist()
+    assert conversations, "The item does not have any conversations."
+    for conversation in item.get("conversations", []):
+        role = conversation["from"]
+        content = conversation["value"]
+        messages.append({"role": role, "content": content})
+
+    return messages
+
+
+def transform_messages(
+    item,
+    frm="chatml",
+    to="text",
+    add_generation_prompt=True,
+    tokenizer=None,
+    assistant_prefix=None,
+):
+    assert to in [
+        "chatml",
+        "text",
+        "sharegpt",
+        "simulated_chat",
+    ], "The output format is not recognized. Please specify the output format."
+    item = deepcopy(item)
+
+    if tokenizer is not None:
+        assert frm == "chatml", "Tokenizer is only supported for chatml format."
+        prompt = tokenizer.apply_chat_template(
+            item, tokenize=False, add_generation_prompt=True
+        )
+        assert isinstance(prompt, str), "Prompt must be a string."
+        if assistant_prefix:
+            prompt += f"{assistant_prefix}"
+        return prompt
+
+    if frm != to:
+        chatml_messages = transform_messages_to_chatml(item, input_format=frm)
+        if to == "sharegpt":
+            if chatml_messages[0]["role"] == "system":
+                system_message = chatml_messages[0]["content"]
+                ret = {"conversations": [], "system": system_message.strip()}
+                for message in chatml_messages[1:]:
+                    ret["conversations"].append(
+                        {"from": message["role"], "value": message["content"]}
+                    )
+            else:
+                system_message = "You are a helpful assistant."
+                ret = {"conversations": [], "system": system_message.strip()}
+                for message in chatml_messages:
+                    ret["conversations"].append(
+                        {"from": message["role"], "value": message["content"]}
+                    )
+            return ret
+        elif to == "chatml":
+            return _transform_sharegpt_to_chatml(item)
+        elif 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":
+            text = "<role> Given the simulated chat, you are the assistant. Lets continue the conversation. \n\n"
+            for turn in chatml_messages:
+                prefix = {
+                    "user": "Human",
+                    "assistant": "AI",
+                    "system": "System",
+                    "function": "Function",
+                }.get(turn["role"])
+                text += f"{prefix}: {turn['content'].strip()}\n\n"
+            if add_generation_prompt:
+                text += "AI: [continue the conversation here]"
+            return text
+        else:
+            raise ValueError(f"{to} is not supported.")
+
+    else:
+        return item
+
+
+def transform_messages_to_chatml(input_data, input_format="auto"):
+    if input_format == "auto":
+        input_data = raw_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."
+        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."
+            input_format = "chatlm"
+            parts = input_data.split("<|im_end|>")
+            input_data = []
+            for part in parts:
+                if not part.strip():
+                    continue
+                role = part.split("<|im_start|>")[1].split("\n")[0]
+                content = part.split(f"<|im_start|>{role}\n")[1]
+                content = content.split("<|im_end|>")[0]
+                input_data.append({"role": role.strip(), "content": content.strip()})
+
+    return input_data
+
+
+__all__ = [
+    "identify_format",
+    "_transform_sharegpt_to_chatml",
+    "transform_messages",
+    "transform_messages_to_chatml",
+]

llm_utils/chat_format/utils.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+from typing import List, Dict, Callable
+
+
+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)
+        kwargs.pop("system_msg", None)
+        for param in params:
+            if param not in kwargs:
+                raise ValueError(f"Missing parameter: {param}")
+        content = template.format(**kwargs)
+        msgs = []
+        if system_msg:
+            msgs += [{"role": "system", "content": system_msg}]
+        msgs += [{"role": "user", "content": content}]
+        return msgs
+
+    return formator
+
+
+def _color_text(text, color_code):
+    return f"\033[{color_code}m{text}\033[0m"
+
+
+def format_msgs(messages):
+    from .transform import transform_messages_to_chatml
+
+    messages = transform_messages_to_chatml(messages)
+    output = []
+    for msg in messages:
+        role = msg.get("role", "unknown").lower()
+        content = msg.get("content", "").strip()
+        output.append(f"{role.capitalize()}:\t{content}")
+        output.append("---")
+    return "\n".join(output)
+
+
+__all__ = [
+    "build_chatml_input",
+    "_color_text",
+    "format_msgs",
+]

llm_utils/group_messages.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import random
+from typing import Sequence, Optional
+
+import numpy as np
+import pandas as pd
+from tabulate import tabulate
+
+from speedy_utils import multi_thread
+
+
+def split_indices_by_length(
+    lengths: Sequence[int],
+    batch_size_by_mean_length: int,
+    random_seed: int,
+    verbose: bool,
+    shuffle: bool,
+    mean_length: Optional[int] = None,
+) -> list[list[int]]:
+    """
+    Split indices into batches so that the sum of lengths in each batch does not exceed max_batch_length.
+    """
+    if mean_length is None:
+        mean_length = int(np.mean(lengths))
+    max_batch_length: int = mean_length * batch_size_by_mean_length
+
+    r: random.Random = random.Random(random_seed)
+    indices: list[int] = list(range(len(lengths)))
+
+    if shuffle:
+        r.shuffle(indices)
+
+    batches: list[list[int]] = []
+    current_batch: list[int] = []
+    current_batch_length: int = 0
+
+    for idx in indices:
+        length: int = lengths[idx]
+        if current_batch_length + length <= max_batch_length:
+            current_batch.append(idx)
+            current_batch_length += length
+        else:
+            batches.append(current_batch)
+            current_batch = [idx]
+            current_batch_length = length
+
+    if current_batch:
+        batches.append(current_batch)
+
+    if verbose:
+        batch_lengths: list[int] = [
+            sum(lengths[idx] for idx in batch) for batch in batches
+        ]
+        desc = pd.Series(batch_lengths).describe()
+
+        table = [
+            ["New avg item len", desc["mean"]],
+            ["Number groups", len(batches)],
+            ["Max length", max_batch_length],
+        ]
+
+        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",
+    batch_size: int = 4,
+    mean_length: int = 512,
+) -> list[dict]:
+    """
+    Groups messages into batches based on token length and concatenates them.
+    """
+    if messages is None:
+        raise ValueError("messages parameter cannot be None")
+    from transformers.models.auto.tokenization_auto import AutoTokenizer
+
+    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)
+            return len(ids)
+
+        lengths: list[int] = multi_thread(get_token_length, messages, workers=64)
+        list_ids: list[list[int]] = split_indices_by_length(
+            lengths,
+            batch_size,
+            random_seed=0,
+            verbose=True,
+            shuffle=True,
+            mean_length=mean_length,
+        )
+        concatenated_messages: list[dict] = []
+
+        def concatenate_messages(conversations: Sequence[Sequence[dict]]) -> dict:
+            system_message = conversations[0][0]
+            turns: list[dict] = []
+            for conv in conversations:
+                turns.extend(conv[1:])
+            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]
+            concatenated_messages.append(concatenate_messages(conversations))
+        return concatenated_messages
+
+    chunked_messages: list[dict] = create_batches(messages)
+    return chunked_messages
+
+
+__all__ = [
+    "split_indices_by_length",
+    "group_messages_by_len",
+]

llm_utils/lm/__init__.py

@@ -0,0 +1,8 @@
+from .lm import LM
+
+OAI_LM = LM
+
+__all__ = [
+    "LM",
+    "OAI_LM",
+]