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

llm_utils/__init__.py

@@ -1,12 +1,13 @@
 from .chat_format import (
     transform_messages,
     transform_messages_to_chatml,
-    display_chat_messages_as_html,
+    show_chat,
     get_conversation_one_turn,
-    display_diff_two_string,
+    show_string_diff,
     display_conversations,
     build_chatml_input,
     format_msgs,
+    display_chat_messages_as_html,
 )
 from .lm import LM
 from .group_messages import (
@@ -17,13 +18,14 @@ from .group_messages import (
 __all__ = [
     "transform_messages",
     "transform_messages_to_chatml",
-    "display_chat_messages_as_html",
+    "show_chat",
     "get_conversation_one_turn",
-    "display_diff_two_string",
+    "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/lm/__init__.py

@@ -1,4 +1,4 @@
-from .base_lm import LM
+from .lm import LM
 
 OAI_LM = LM
 

llm_utils/lm/{base_lm.py → lm.py}

@@ -60,7 +60,7 @@ class LM:
         temperature: float = 0.0,
         max_tokens: int = 2_000,
         host: str = "localhost",
-        port: Optional[int] = None,
+        port: Optional[int | str] = None,
         base_url: Optional[str] = None,
         api_key: Optional[str] = None,
         cache: bool = True,
@@ -289,12 +289,12 @@ class LM:
             return None
 
     @staticmethod
-    def list_models(port=None) -> List[str]:
+    def list_models(port=None, host='localhost') -> List[str]:
         """
         List available models.
         """
         try:
-            client: OpenAI = LM(port=port).client
+            client: OpenAI = LM(port=port, host=host).client
             base_url: URL = client.base_url
             logger.debug(f"Base URL: {base_url}")
             models: SyncPage[Model] = client.models.list()