speedy-utils 1.0.15__py3-none-any.whl → 1.0.20__py3-none-any.whl

llm_utils/__init__.py

@@ -9,8 +9,8 @@ from .chat_format import (
     transform_messages,
     transform_messages_to_chatml,
 )
-from .lm.lm import LM, LLMTask
-from .lm.alm import AsyncLM
+from .lm.async_lm import AsyncLLMTask, AsyncLM
+from .lm.sync_lm import LM, LLMTask
 
 __all__ = [
     "transform_messages",
@@ -26,4 +26,5 @@ __all__ = [
     "AsyncLM",
     "display_chat_messages_as_html",
     "LLMTask",
+    "AsyncLLMTask",
 ]

llm_utils/chat_format/display.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import List, Tuple, Sequence, Any, Dict, Optional
+from typing import Any, Optional
 from IPython.display import HTML, display
 from difflib import SequenceMatcher
 

llm_utils/chat_format/transform.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 from copy import deepcopy
-from typing import Callable, Dict, List, Sequence
 
 
 def identify_format(item):
@@ -114,7 +113,7 @@ def transform_messages(
 
 def transform_messages_to_chatml(input_data, input_format="auto"):
     if input_format == "auto":
-        input_data = raw_data = deepcopy(input_data)
+        input_data = deepcopy(input_data)
         if isinstance(input_data, list):
             input_format = "chatlm"
             assert (

llm_utils/group_messages.py

@@ -76,7 +76,7 @@ def group_messages_by_len(
     """
     if messages is None:
         raise ValueError("messages parameter cannot be None")
-    from transformers.models.auto.tokenization_auto import AutoTokenizer
+    from transformers.models.auto.tokenization_auto import AutoTokenizer # type: ignore
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 

llm_utils/lm/__init__.py

@@ -1,9 +1,12 @@
-from .lm import LM
-from .alm import AsyncLM
-OAI_LM = LM
+# from .async_lm import AsyncLLMTask, AsyncLM
+# from .lm import LM, LLMTask
 
-__all__ = [
-    "LM",
-    "OAI_LM",
-    "AsyncLM",
-]
+# OAI_LM = LM
+
+# __all__ = [
+#     "LM",
+#     "OAI_LM",
+#     "AsyncLM",
+#     "LLMTask",
+#     "AsyncLLMTask",
+# ]

llm_utils/lm/{alm.py → async_lm.py}

@@ -1,5 +1,3 @@
-from __future__ import annotations
-
 """An **asynchronous** drop‑in replacement for the original `LM` class.
 
 Usage example (Python ≥3.8):
@@ -15,26 +13,30 @@ Usage example (Python ≥3.8):
     asyncio.run(main())
 """
 
-import asyncio
 import base64
 import hashlib
 import json
 import os
+from abc import ABC
+from functools import lru_cache
 from typing import (
     Any,
     Dict,
     List,
+    Literal,
     Optional,
     Sequence,
     Type,
     TypeVar,
     Union,
-    overload,
     cast,
+    overload,
 )
 
 from httpx import URL
+from loguru import logger
 from openai import AsyncOpenAI, AuthenticationError, BadRequestError, RateLimitError
+from openai.pagination import AsyncPage as AsyncSyncPage
 
 # from openai.pagination import AsyncSyncPage
 from openai.types.chat import (
@@ -44,11 +46,10 @@ from openai.types.chat import (
     ChatCompletionToolMessageParam,
     ChatCompletionUserMessageParam,
 )
-from openai.types.chat.parsed_chat_completion import ParsedChatCompletion
 from openai.types.model import Model
 from pydantic import BaseModel
-from loguru import logger
-from openai.pagination import AsyncPage as AsyncSyncPage
+
+from llm_utils.chat_format.display import get_conversation_one_turn
 
 # --------------------------------------------------------------------------- #
 # type helpers
@@ -67,10 +68,20 @@ def _color(code: int, text: str) -> str:
     return f"\x1b[{code}m{text}\x1b[0m"
 
 
-_red = lambda t: _color(31, t)
-_green = lambda t: _color(32, t)
-_blue = lambda t: _color(34, t)
-_yellow = lambda t: _color(33, t)
+def _red(t):
+    return _color(31, t)
+
+
+def _green(t):
+    return _color(32, t)
+
+
+def _blue(t):
+    return _color(34, t)
+
+
+def _yellow(t):
+    return _color(33, t)
 
 
 class AsyncLM:
@@ -100,6 +111,7 @@ class AsyncLM:
         self.openai_kwargs = openai_kwargs
         self.do_cache = cache
         self.ports = ports
+        self._init_port = port  # <-- store the port provided at init
 
         # Async client
 
@@ -375,10 +387,182 @@ class AsyncLM:
         except Exception:
             return None
 
+    # ------------------------------------------------------------------ #
+    # Missing methods from LM class
+    # ------------------------------------------------------------------ #
+    async def parse(
+        self,
+        response_model: Type[BaseModel],
+        instruction: Optional[str] = None,
+        prompt: Optional[str] = None,
+        messages: Optional[RawMsgs] = None,
+        think: Literal[True, False, None] = None,
+        add_json_schema_to_instruction: bool = False,
+        temperature: Optional[float] = None,
+        max_tokens: Optional[int] = None,
+        return_openai_response: bool = False,
+        cache: Optional[bool] = True,
+        **kwargs,
+    ):
+        """Parse response using guided JSON generation."""
+        if messages is None:
+            assert instruction is not None, "Instruction must be provided."
+            assert prompt is not None, "Prompt must be provided."
+            messages = [
+                {
+                    "role": "system",
+                    "content": instruction,
+                },
+                {
+                    "role": "user",
+                    "content": prompt,
+                },
+            ]  # type: ignore
+
+        post_fix = ""
+        json_schema = response_model.model_json_schema()
+        if add_json_schema_to_instruction and response_model:
+            _schema = f"\n\n<output_json_schema>\n{json.dumps(json_schema, indent=2)}\n</output_json_schema>"
+            post_fix += _schema
+
+        if think:
+            post_fix += "\n\n/think"
+        elif not think:
+            post_fix += "\n\n/no_think"
+
+        assert isinstance(messages, list), "Messages must be a list."
+        assert len(messages) > 0, "Messages cannot be empty."
+        assert messages[0]["role"] == "system", (
+            "First message must be a system message with instruction."
+        )
+        messages[0]["content"] += post_fix  # type: ignore
+
+        model_kwargs = {}
+        if temperature is not None:
+            model_kwargs["temperature"] = temperature
+        if max_tokens is not None:
+            model_kwargs["max_tokens"] = max_tokens
+        model_kwargs.update(kwargs)
+
+        use_cache = self.do_cache if cache is None else cache
+        cache_key = None
+        if use_cache:
+            cache_data = {
+                "messages": messages,
+                "model_kwargs": model_kwargs,
+                "guided_json": json_schema,
+                "response_format": response_model.__name__,
+            }
+            cache_key = self._cache_key(cache_data, {}, response_model)
+            cached_response = self._load_cache(cache_key)
+            self.last_log = [prompt, messages, cached_response]
+            if cached_response is not None:
+                if return_openai_response:
+                    return cached_response
+                return self._parse_complete_output(cached_response, response_model)
+
+        completion = await self.client.chat.completions.create(
+            model=self.model,  # type: ignore
+            messages=messages,  # type: ignore
+            extra_body={"guided_json": json_schema},
+            **model_kwargs,
+        )
+
+        if cache_key:
+            self._dump_cache(cache_key, completion)
+
+        self.last_log = [prompt, messages, completion]
+        if return_openai_response:
+            return completion
+        return self._parse_complete_output(completion, response_model)
+
+    def _parse_complete_output(
+        self, completion: Any, response_model: Type[BaseModel]
+    ) -> BaseModel:
+        """Parse completion output to response model."""
+        if hasattr(completion, "model_dump"):
+            completion = completion.model_dump()
+
+        if "choices" not in completion or not completion["choices"]:
+            raise ValueError("No choices in OpenAI response")
+
+        content = completion["choices"][0]["message"]["content"]
+        if not content:
+            raise ValueError("Empty content in response")
+
+        try:
+            data = json.loads(content)
+            return response_model.model_validate(data)
+        except Exception as exc:
+            raise ValueError(
+                f"Failed to parse response as {response_model.__name__}: {content}"
+            ) from exc
+
+    async def inspect_word_probs(
+        self,
+        messages: Optional[List[Dict[str, Any]]] = None,
+        tokenizer: Optional[Any] = None,
+        do_print=True,
+        add_think: bool = True,
+    ) -> tuple[List[Dict[str, Any]], Any, str]:
+        """
+        Inspect word probabilities in a language model response.
+
+        Args:
+            tokenizer: Tokenizer instance to encode words.
+            messages: List of messages to analyze.
+
+        Returns:
+            A tuple containing:
+            - List of word probabilities with their log probabilities.
+            - Token log probability dictionaries.
+            - Rendered string with colored word probabilities.
+        """
+        if messages is None:
+            messages = await self.last_messages(add_think=add_think)
+            if messages is None:
+                raise ValueError("No messages provided and no last messages available.")
+
+        if tokenizer is None:
+            tokenizer = get_tokenizer(self.model)
+
+        ret = await inspect_word_probs_async(self, tokenizer, messages)
+        if do_print:
+            print(ret[-1])
+        return ret
+
+    async def last_messages(
+        self, add_think: bool = True
+    ) -> Optional[List[Dict[str, str]]]:
+        """Get the last conversation messages including assistant response."""
+        if not hasattr(self, "last_log"):
+            return None
+
+        last_conv = self.last_log
+        messages = last_conv[1] if len(last_conv) > 1 else None
+        last_msg = last_conv[2]
+        if not isinstance(last_msg, dict):
+            last_conv[2] = last_conv[2].model_dump()  # type: ignore
+        msg = last_conv[2]
+        # Ensure msg is a dict
+        if hasattr(msg, "model_dump"):
+            msg = msg.model_dump()
+        message = msg["choices"][0]["message"]
+        reasoning = message.get("reasoning_content")
+        answer = message.get("content")
+        if reasoning and add_think:
+            final_answer = f"<think>{reasoning}</think>\n{answer}"
+        else:
+            final_answer = f"<think>\n\n</think>\n{answer}"
+        assistant = {"role": "assistant", "content": final_answer}
+        messages = messages + [assistant]  # type: ignore
+        return messages if messages else None
+
     # ------------------------------------------------------------------ #
     # Utility helpers
     # ------------------------------------------------------------------ #
     async def inspect_history(self) -> None:
+        """Inspect the conversation history with proper formatting."""
         if not hasattr(self, "last_log"):
             raise ValueError("No history available. Please call the model first.")
 
@@ -466,3 +650,210 @@ class AsyncLM:
         except Exception as exc:
             logger.error(f"Failed to list models: {exc}")
             return []
+
+
+# --------------------------------------------------------------------------- #
+# Module-level utility functions (async versions)
+# --------------------------------------------------------------------------- #
+
+
+@lru_cache(maxsize=10)
+def get_tokenizer(model_name: str) -> Any:
+    """Get tokenizer for the given model."""
+    from transformers import AutoTokenizer  # type: ignore
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    return tokenizer
+
+
+async def inspect_word_probs_async(lm, tokenizer, messages):
+    """Async version of inspect_word_probs."""
+
+    import numpy as np
+
+    async def compute_word_log_probs(
+        tokenizer: Any,
+        lm_client: Any,
+    ) -> tuple[List[Dict[str, Any]], Any]:
+        # Build a prompt that preserves literal newlines
+        prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,  # Don't tokenize yet, we need raw text
+            add_generation_prompt=False,  # No generation prompt needed
+        )
+
+        # Request token logprobs
+        response = await lm_client.client.completions.create(
+            model=lm_client.model,  # type: ignore
+            prompt=prompt,
+            max_tokens=1,
+            logprobs=1,
+            extra_body={"prompt_logprobs": 0},
+        )
+        token_logprob_dicts = response.choices[0].prompt_logprobs  # type: ignore
+
+        # Override first token to known start marker
+        start_id = tokenizer.encode("<|im_start|>")[0]
+        token_logprob_dicts[0] = {
+            str(start_id): {
+                "logprob": -1,
+                "rank": 1,
+                "decoded_token": "<|im_start|>",
+            }
+        }
+
+        # Flatten tokens
+        tokens: List[Dict[str, Any]] = [
+            {"id": int(tid), **tdata}
+            for td in token_logprob_dicts
+            for tid, tdata in td.items()
+        ]
+
+        # Validate tokenization
+        tokenized = tokenizer.tokenize(prompt)
+        if len(tokenized) != len(tokens):
+            raise ValueError(f"Token count mismatch: {len(tokenized)} vs {len(tokens)}")
+        for idx, tok in enumerate(tokens):
+            if tokenized[idx] != tok["decoded_token"]:
+                raise AssertionError(
+                    f"Token mismatch at {idx}: "
+                    f"{tokenized[idx]} != {tok['decoded_token']}"
+                )
+
+        # Split on newline sentinel
+        split_prompt = prompt.replace("\n", " <NL> ")
+        words = split_prompt.split()
+
+        word_log_probs: List[Dict[str, Any]] = []
+        token_idx = 0
+
+        for word in words:
+            # Map sentinel back to actual newline for encoding
+            target = "\n" if word == "<NL>" else word
+            sub_ids = tokenizer.encode(target, add_special_tokens=False)
+            count = len(sub_ids)
+            if count == 0:
+                continue
+
+            subs = tokens[token_idx : token_idx + count]
+            avg_logprob = sum(s["logprob"] for s in subs) / count
+            prob = float(np.exp(avg_logprob))
+            word_log_probs.append({"word": target, "probability": prob})
+            token_idx += count
+
+        return word_log_probs, token_logprob_dicts  # type: ignore
+
+    def render_by_logprob(word_log_probs: List[Dict[str, Any]]) -> str:
+        """
+        Return an ANSI-colored string for word probabilities (red → green).
+        """
+        if not word_log_probs:
+            return ""
+
+        probs = [entry["probability"] for entry in word_log_probs]
+        min_p, max_p = min(probs), max(probs)
+        parts: List[str] = []
+
+        for entry in word_log_probs:
+            word = entry["word"]
+            # Preserve actual line breaks
+            if word == "\n":
+                parts.append("\n")
+                continue
+
+            p = entry["probability"]
+            norm = (p - min_p) / (max_p - min_p or 1.0)
+            r = int(255 * (1 - norm))  # red component (high when prob is low)
+            g = int(255 * norm)  # green component (high when prob is high)
+            b = 0  # no blue for red-green gradient
+            colored = f"\x1b[38;2;{r};{g};{b}m{word}\x1b[0m"
+            parts.append(colored + " ")
+
+        return "".join(parts).rstrip()
+
+    word_probs, token_logprob_dicts = await compute_word_log_probs(tokenizer, lm)
+    return word_probs, token_logprob_dicts, render_by_logprob(word_probs)
+
+
+# --------------------------------------------------------------------------- #
+# Async LLMTask class
+# --------------------------------------------------------------------------- #
+
+
+class AsyncLLMTask(ABC):
+    """
+    Async callable wrapper around an AsyncLM endpoint.
+
+    Sub-classes must set:
+      • lm              – the async language-model instance
+      • InputModel      – a Pydantic input class
+      • OutputModel     – a Pydantic output class
+
+    Optional flags:
+      • temperature     – float (default 0.6)
+      • think           – bool  (if the backend supports "chain-of-thought")
+      • add_json_schema – bool  (include schema in the instruction)
+
+    The **docstring** of each sub-class is sent as the LM instruction.
+    Example
+    ```python
+        class DemoTask(AsyncLLMTask):
+            "TODO: SYSTEM_PROMPT_INSTURCTION HERE"
+
+            lm = AsyncLM(port=8130, cache=False, model="gpt-3.5-turbo")
+
+            class InputModel(BaseModel):
+                text_to_translate:str
+
+            class OutputModel(BaseModel):
+                translation:str
+                glossary_use:str
+
+            temperature = 0.6
+            think=False
+
+        demo_task = DemoTask()
+        result = await demo_task({'text_to_translate': 'Translate from english to vietnamese: Hello how are you'})
+    ```
+    """
+
+    lm: "AsyncLM"
+    InputModel: Type[BaseModel]
+    OutputModel: Type[BaseModel]
+
+    temperature: float = 0.6
+    think: bool = False
+    add_json_schema: bool = False
+
+    async def __call__(self, data: BaseModel | dict) -> BaseModel:
+        if (
+            not hasattr(self, "InputModel")
+            or not hasattr(self, "OutputModel")
+            or not hasattr(self, "lm")
+        ):
+            raise NotImplementedError(
+                f"{self.__class__.__name__} must define lm, InputModel, and OutputModel as class attributes."
+            )
+
+        item = data if isinstance(data, BaseModel) else self.InputModel(**data)
+
+        return await self.lm.parse(
+            prompt=item.model_dump_json(),
+            instruction=self.__doc__ or "",
+            response_model=self.OutputModel,
+            temperature=self.temperature,
+            think=self.think,
+            add_json_schema_to_instruction=self.add_json_schema,
+        )
+
+    def generate_training_data(
+        self, input_dict: Dict[str, Any], output: Dict[str, Any]
+    ):
+        """Return share gpt like format"""
+        system_prompt = self.__doc__ or ""
+        user_msg = self.InputModel(**input_dict).model_dump_json()  # type: ignore[attr-defined]
+        assistant_msg = self.OutputModel(**output).model_dump_json()  # type: ignore[attr-defined]
+        messages = get_conversation_one_turn(
+            system_msg=system_prompt, user_msg=user_msg, assistant_msg=assistant_msg
+        )
+        return {"messages": messages}

llm_utils/lm/chat_html.py

@@ -1,4 +1,5 @@
-from .lm import *
+from typing import Any, Optional, cast
+from .sync_lm import LM, Messages, LegacyMsgs, RawMsgs
 import sys
 
 # Configuration
@@ -231,7 +232,8 @@ class LMChatHtml(LM):
                         padding=display_padding,
                         inner_padding=display_inner_padding,
                     )
-                    display_handle.update(HTML(html_content))
+                    if display_handle is not None:
+                        display_handle.update(HTML(html_content))
         else:
             # Console streaming mode (original behavior)
             for chunk in stream:

llm_utils/lm/lm_json.py

@@ -1,11 +1,7 @@
-import json
-import re
-from functools import cache
-from typing import *
+from typing import Any, Optional
 
-from pydantic import BaseModel
 
-from llm_utils.lm.lm import LM, RawMsgs
+from llm_utils.lm.sync_lm import LM, RawMsgs
 
 
 class LMJson(LM):

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

@@ -4,8 +4,8 @@ import base64
 import hashlib
 import json
 import os
-import warnings
 from abc import ABC
+from functools import lru_cache
 from typing import (
     Any,
     Dict,
@@ -20,10 +20,7 @@ from typing import (
     overload,
 )
 
-from httpx import URL
-from huggingface_hub import repo_info
 from loguru import logger
-from numpy import isin
 from openai import AuthenticationError, OpenAI, RateLimitError
 from openai.pagination import SyncPage
 from openai.types.chat import (
@@ -33,10 +30,10 @@ from openai.types.chat import (
     ChatCompletionToolMessageParam,
     ChatCompletionUserMessageParam,
 )
-from openai.types.chat.parsed_chat_completion import ParsedChatCompletion
 from openai.types.model import Model
 from pydantic import BaseModel
 
+from llm_utils.chat_format.display import get_conversation_one_turn
 from speedy_utils.common.utils_io import jdumps
 
 # --------------------------------------------------------------------------- #
@@ -549,7 +546,7 @@ class LM:
 
         if think:
             post_fix += "\n\n/think"
-        elif think == False:
+        elif not think:
             post_fix += "\n\n/no_think"
 
         assert isinstance(messages, list), "Messages must be a list."
@@ -675,20 +672,15 @@ class LM:
         return messages if messages else None
 
 
-from functools import lru_cache
-
-
 @lru_cache(maxsize=10)
 def get_tokenizer(model_name: str) -> Any:
-    from transformers import AutoTokenizer
+    from transformers import AutoTokenizer # type: ignore
 
     tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
     return tokenizer
 
 
 def inspect_word_probs(lm, tokenizer, messages):
-    import re
-    from typing import Any, Dict, List
 
     import numpy as np
 

llm_utils/lm/utils.py

@@ -82,7 +82,7 @@ def retry_on_exception(max_retries=10, exceptions=(Exception,), sleep_time=3):
                 try:
                     return func(self, *args, **kwargs)
                 except exceptions as e:
-                    import litellm
+                    import litellm # type: ignore
 
                     if isinstance(
                         e, (litellm.exceptions.APIError, litellm.exceptions.Timeout)

llm_utils/scripts/vllm_load_balancer.py

@@ -1,7 +1,6 @@
 import asyncio
 import random
 from collections import defaultdict
-import time
 from tabulate import tabulate
 import contextlib
 import aiohttp  # <-- Import aiohttp

speedy_utils/common/function_decorator.py

@@ -2,7 +2,7 @@ import functools
 import time
 import traceback
 from collections.abc import Callable
-from typing import Any, Tuple, Type, Union
+from typing import Any
 
 from loguru import logger
 
@@ -26,8 +26,6 @@ def retry_runtime(
     def decorator(func: Callable) -> Callable:
         @functools.wraps(func)
         def wrapper(*args: Any, **kwargs: Any) -> Any:
-            last_exception = None
-
             for attempt in range(1, max_retry + 1):
                 try:
                     return func(*args, **kwargs)
@@ -40,7 +38,6 @@ def retry_runtime(
                     raise
 
                 except exceptions as e:
-                    last_exception = e
                     if attempt == max_retry:
                         logger.opt(depth=1).error(
                             f"Function {func.__name__} failed after {max_retry} retries: {str(e)}"

speedy_utils/common/logger.py

@@ -5,7 +5,7 @@ import re
 import sys
 import time
 from collections import OrderedDict
-from typing import Annotated, Literal, Optional
+from typing import Annotated, Literal
 
 from loguru import logger
 

speedy_utils/common/report_manager.py

@@ -2,7 +2,6 @@ import os
 from collections import defaultdict
 from datetime import datetime
 
-from fastcore.all import threaded
 
 
 class ReportManager:
@@ -40,7 +39,7 @@ class ReportManager:
 
         md_content.extend(
             [
-                f"\n### Results Overview",
+                "\n### Results Overview",
                 f"- Total items processed: {len(results)}",
                 f"- Success rate: {(len(results) - len(errors))/len(results)*100:.1f}%",
                 f"- Total errors: {len(errors)}",
@@ -61,7 +60,7 @@ class ReportManager:
             for error_type, errs in error_groups.items():
                 md_content.extend(
                     [
-                        f"\n<details>",
+                        "\n<details>",
                         f"<summary><b>{error_type}</b> ({len(errs)} occurrences)</summary>\n",
                         "| Index | Input | Error Message |",
                         "|-------|-------|---------------|",

speedy_utils/common/utils_cache.py

@@ -5,7 +5,8 @@ import os
 import os.path as osp
 import pickle
 import uuid
-from typing import Any, List, Literal
+from threading import Lock
+from typing import Any, Literal
 
 import cachetools
 import pandas as pd
@@ -13,12 +14,11 @@ import xxhash
 from loguru import logger
 from pydantic import BaseModel
 
-from .utils_io import dump_json_or_pickle, load_json_or_pickle
-from .utils_misc import mkdir_or_exist
+from src.speedy_utils.common.utils_io import dump_json_or_pickle, load_json_or_pickle
+from src.speedy_utils.common.utils_misc import mkdir_or_exist
 
 SPEED_CACHE_DIR = osp.join(osp.expanduser("~"), ".cache/speedy_cache")
 LRU_MEM_CACHE = cachetools.LRUCache(maxsize=128_000)
-from threading import Lock
 
 thread_locker = Lock()
 

speedy_utils/common/utils_misc.py

@@ -2,9 +2,8 @@
 
 import inspect
 import os
-import sys
 from collections.abc import Callable
-from typing import Any, List
+from typing import Any
 
 from pydantic import BaseModel
 

speedy_utils/multi_worker/process.py

@@ -1,12 +1,10 @@
-import inspect
 import multiprocessing
 import os
-import time
 import traceback
-from collections.abc import Callable, Iterable, Iterator, Sequence
+from collections.abc import Callable, Iterable, Iterator
 from concurrent.futures import ProcessPoolExecutor, as_completed
 from itertools import islice
-from typing import Any, List, TypeVar, cast
+from typing import Any, TypeVar, cast
 
 T = TypeVar("T")
 

speedy_utils/scripts/mpython.py

@@ -5,7 +5,6 @@ import multiprocessing  # Import multiprocessing module
 import os
 import shlex  # To properly escape command line arguments
 import shutil
-import subprocess
 
 taskset_path = shutil.which("taskset")
 
@@ -80,7 +79,7 @@ def main():
         cmd_str = shlex.join(args.cmd)
 
     gpus = args.gpus.split(",")
-    gpus = [gpu for gpu in gpus if not gpu in args.ignore_gpus.split(",")]
+    gpus = [gpu for gpu in gpus if gpu not in args.ignore_gpus.split(",")]
     num_gpus = len(gpus)
 
     cpu_per_process = max(args.total_cpu // args.total_fold, 1)

speedy_utils/scripts/openapi_client_codegen.py

@@ -10,12 +10,8 @@ Usage:
 
 import argparse
 import json
-import os
-import signal
-import subprocess
 import sys
-import time
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List
 
 
 def pascal_case(s: str) -> str:

{speedy_utils-1.0.15.dist-info → speedy_utils-1.0.20.dist-info}/METADATA

@@ -1,11 +1,12 @@
 Metadata-Version: 2.3
 Name: speedy-utils
-Version: 1.0.15
+Version: 1.0.20
 Summary: Fast and easy-to-use package for data science
 Author: AnhVTH
 Author-email: anhvth.226@gmail.com
-Requires-Python: >=3.9
+Requires-Python: >=3.8
 Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
@@ -19,7 +20,7 @@ Requires-Dist: fastprogress
 Requires-Dist: freezegun (>=1.5.1,<2.0.0)
 Requires-Dist: ipdb
 Requires-Dist: ipywidgets
-Requires-Dist: json-repair (>=0.40.0,<0.41.0)
+Requires-Dist: json-repair (>=0.25.0,<0.31.0)
 Requires-Dist: jupyterlab
 Requires-Dist: loguru
 Requires-Dist: matplotlib

speedy_utils-1.0.20.dist-info/RECORD

@@ -0,0 +1,37 @@
+llm_utils/__init__.py,sha256=SlaCMArn_uKVw4r0psz0q0IOQ1VFGdgCFOAKxQ81WTI,694
+llm_utils/chat_format/__init__.py,sha256=8dBIUqFJvkgQYedxBtcyxt-4tt8JxAKVap2JlTXmgaM,737
+llm_utils/chat_format/display.py,sha256=qaEGADGP8iQFzWOuzEO7_HyrqAFdEnUfkHAH28b0ym0,9772
+llm_utils/chat_format/transform.py,sha256=8TZhvUS5DrjUeMNtDIuWY54B_QZ7jjpXEL9c8F5z79w,5400
+llm_utils/chat_format/utils.py,sha256=xTxN4HrLHcRO2PfCTR43nH1M5zCa7v0kTTdzAcGkZg0,1229
+llm_utils/group_messages.py,sha256=8CU9nKOja3xeuhdrX5CvYVveSqSKb2zQ0eeNzA88aTQ,3621
+llm_utils/lm/__init__.py,sha256=rX36_MsnekM5GHwWS56XELbm4W5x2TDwnPERDTfo0eU,194
+llm_utils/lm/async_lm.py,sha256=06QVT7iWAa8Rz48oRmnzkS-dJYOiJ0zSfKqEhZY11S8,30825
+llm_utils/lm/chat_html.py,sha256=FkGo0Dv_nAHYBMZzXfMu_bGQKaCx302goh3XaT-_ETc,8674
+llm_utils/lm/lm_json.py,sha256=fMt42phzFV2f6ulrtWcDXsWHi8WcG7gGkCzpIq8VSSM,1975
+llm_utils/lm/sync_lm.py,sha256=FMiAhltqUNZdzGejMBPqGfQ3PE-fUoIGW8hfjTO_fNo,31035
+llm_utils/lm/utils.py,sha256=GMvs64DRzVnXAki4SZ-A6mx2Fi9IgeF11BA-5FB-CYg,4777
+llm_utils/scripts/README.md,sha256=yuOLnLa2od2jp4wVy3rV0rESeiV3o8zol5MNMsZx0DY,999
+llm_utils/scripts/vllm_load_balancer.py,sha256=qsVhQ3ubJ3JnOAu4psHB1PyKy9muX-mZZkFqepH8WcU,17507
+llm_utils/scripts/vllm_serve.py,sha256=CbW_3Y9Vt7eQYoGGPT3yj1nhbLYOc3b1LdJBy1sVX-Y,11976
+speedy_utils/__init__.py,sha256=kxQk4PGS3Xkxnerm0YqjF6GKTpgoaTc1vudKid-2c_A,3388
+speedy_utils/all.py,sha256=A9jiKGjo950eg1pscS9x38OWAjKGyusoAN5mrfweY4E,3090
+speedy_utils/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+speedy_utils/common/clock.py,sha256=3n4FkCW0dz46O8By09V5Pve1DSMgpLDRbWEVRryryeQ,7423
+speedy_utils/common/function_decorator.py,sha256=BspJ0YuGL6elS7lWBAgELZ-sCfED_1N2P5fgH-fCRUQ,2132
+speedy_utils/common/logger.py,sha256=JqW9gG4ujfq4RldNeYP2p52BYgCwjkYeGGYyzLn6mfY,6422
+speedy_utils/common/notebook_utils.py,sha256=7tFXMaE365ubrD9VCwDoZe7oBqZ2qG5t_Vla4Qvg1M8,2062
+speedy_utils/common/report_manager.py,sha256=sEuwLcyYpYlj9uxIsC8yiAu1eirydkQIldKRtFStswI,3894
+speedy_utils/common/utils_cache.py,sha256=Fp1vvi19q4xIcXi7ZuVjk49gOqdl3i7bv2Z6UnERMjM,8424
+speedy_utils/common/utils_io.py,sha256=d7PKz5tOPrwHDr7GXuYmILvjXJOFEwfzAEIuUcYaI60,4790
+speedy_utils/common/utils_misc.py,sha256=cdEuBBpiB1xpuzj0UBDHDuTIerqsMIw37ENq6EXliOw,1795
+speedy_utils/common/utils_print.py,sha256=iQqnOYw2EFC8TqeSDbrcnIQAUKT7FbB8Mec8b2aGAzw,4833
+speedy_utils/multi_worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+speedy_utils/multi_worker/process.py,sha256=BI-sgzzQ0_N8kOfaS_3ZAGZ3d6panYzJ3-BGZthY4dQ,6824
+speedy_utils/multi_worker/thread.py,sha256=9pXjvgjD0s0Hp0cZ6I3M0ndp1OlYZ1yvqbs_bcun_Kw,12775
+speedy_utils/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+speedy_utils/scripts/mpython.py,sha256=73PHm1jqbCt2APN4xuNjD0VDKwzOj4EZsViEMQiZU2g,3853
+speedy_utils/scripts/openapi_client_codegen.py,sha256=f2125S_q0PILgH5dyzoKRz7pIvNEjCkzpi4Q4pPFRZE,9683
+speedy_utils-1.0.20.dist-info/METADATA,sha256=ujYqU8i_7O40c4av_5M16eoPIDsA20df2uOSFUVOY-Q,7442
+speedy_utils-1.0.20.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+speedy_utils-1.0.20.dist-info/entry_points.txt,sha256=T1t85jwx8fK6m5msdkBGIXH5R5Kd0zSL0S6erXERPzg,237
+speedy_utils-1.0.20.dist-info/RECORD,,

llm_utils/scripts/example_vllm_client.py

@@ -1,269 +0,0 @@
-"""
-Beautiful example script for interacting with VLLM server.
-
-This script demonstrates various ways to use the VLLM API server
-for text generation tasks.
-"""
-
-import asyncio
-import json
-from typing import Dict, List, Optional, Any
-
-import aiohttp
-from loguru import logger
-from pydantic import BaseModel, Field
-
-
-class VLLMRequest(BaseModel):
-    """Request model for VLLM API."""
-    prompt: str
-    max_tokens: int = Field(default=512, ge=1, le=8192)
-    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
-    top_p: float = Field(default=0.9, ge=0.0, le=1.0)
-    stream: bool = False
-    stop: Optional[List[str]] = None
-
-
-class VLLMResponse(BaseModel):
-    """Response model from VLLM API."""
-    text: str
-    finish_reason: str
-    prompt_tokens: int
-    completion_tokens: int
-    total_tokens: int
-
-
-class VLLMClient:
-    """Client for interacting with VLLM server."""
-    
-    def __init__(self, base_url: str = 'http://localhost:8140'):
-        self.base_url = base_url
-        self.model_name = 'selfeval_8b'
-        
-    async def generate_text(
-        self, 
-        request: VLLMRequest
-    ) -> VLLMResponse:
-        """Generate text using VLLM API."""
-        url = f'{self.base_url}/v1/completions'
-        
-        payload = {
-            'model': self.model_name,
-            'prompt': request.prompt,
-            'max_tokens': request.max_tokens,
-            'temperature': request.temperature,
-            'top_p': request.top_p,
-            'stream': request.stream,
-        }
-        
-        if request.stop:
-            payload['stop'] = request.stop
-            
-        async with aiohttp.ClientSession() as session:
-            try:
-                async with session.post(
-                    url, 
-                    json=payload,
-                    timeout=aiohttp.ClientTimeout(total=60)
-                ) as response:
-                    response.raise_for_status()
-                    data = await response.json()
-                    
-                    choice = data['choices'][0]
-                    usage = data['usage']
-                    
-                    return VLLMResponse(
-                        text=choice['text'],
-                        finish_reason=choice['finish_reason'],
-                        prompt_tokens=usage['prompt_tokens'],
-                        completion_tokens=usage['completion_tokens'],
-                        total_tokens=usage['total_tokens']
-                    )
-                    
-            except aiohttp.ClientError as e:
-                logger.error(f'HTTP error: {e}')
-                raise
-            except Exception as e:
-                logger.error(f'Unexpected error: {e}')
-                raise
-                
-    async def generate_batch(
-        self, 
-        requests: List[VLLMRequest]
-    ) -> List[VLLMResponse]:
-        """Generate text for multiple requests concurrently."""
-        tasks = [self.generate_text(req) for req in requests]
-        return await asyncio.gather(*tasks, return_exceptions=True)
-        
-    async def health_check(self) -> bool:
-        """Check if the VLLM server is healthy."""
-        url = f'{self.base_url}/health'
-        
-        try:
-            async with aiohttp.ClientSession() as session:
-                async with session.get(
-                    url,
-                    timeout=aiohttp.ClientTimeout(total=10)
-                ) as response:
-                    return response.status == 200
-        except Exception as e:
-            logger.warning(f'Health check failed: {e}')
-            return False
-
-
-async def example_basic_generation():
-    """Example: Basic text generation."""
-    logger.info('🚀 Running basic generation example')
-    
-    client = VLLMClient()
-    
-    # Check server health
-    if not await client.health_check():
-        logger.error('❌ Server is not healthy')
-        return
-        
-    request = VLLMRequest(
-        prompt='Explain the concept of machine learning in simple terms:',
-        max_tokens=256,
-        temperature=0.7,
-        stop=['\n\n']
-    )
-    
-    try:
-        response = await client.generate_text(request)
-        
-        logger.success('✅ Generation completed')
-        logger.info(f'📝 Generated text:\n{response.text}')
-        logger.info(f'📊 Tokens: {response.total_tokens} total '
-                   f'({response.prompt_tokens} prompt + '
-                   f'{response.completion_tokens} completion)')
-                   
-    except Exception as e:
-        logger.error(f'❌ Generation failed: {e}')
-
-
-async def example_batch_generation():
-    """Example: Batch text generation."""
-    logger.info('🚀 Running batch generation example')
-    
-    client = VLLMClient()
-    
-    prompts = [
-        'What is artificial intelligence?',
-        'Explain quantum computing briefly:',
-        'What are the benefits of renewable energy?'
-    ]
-    
-    requests = [
-        VLLMRequest(
-            prompt=prompt,
-            max_tokens=128,
-            temperature=0.8
-        ) for prompt in prompts
-    ]
-    
-    try:
-        responses = await client.generate_batch(requests)
-        
-        for i, response in enumerate(responses):
-            if isinstance(response, Exception):
-                logger.error(f'❌ Request {i+1} failed: {response}')
-            else:
-                logger.success(f'✅ Request {i+1} completed')
-                logger.info(f'📝 Response {i+1}:\n{response.text}\n')
-                
-    except Exception as e:
-        logger.error(f'❌ Batch generation failed: {e}')
-
-
-async def example_creative_writing():
-    """Example: Creative writing with specific parameters."""
-    logger.info('🚀 Running creative writing example')
-    
-    client = VLLMClient()
-    
-    request = VLLMRequest(
-        prompt=(
-            'Write a short story about a robot discovering emotions. '
-            'The story should be exactly 3 paragraphs:\n\n'
-        ),
-        max_tokens=400,
-        temperature=1.2,  # Higher temperature for creativity
-        top_p=0.95,
-        stop=['THE END', '\n\n\n']
-    )
-    
-    try:
-        response = await client.generate_text(request)
-        
-        logger.success('✅ Creative writing completed')
-        logger.info(f'📚 Story:\n{response.text}')
-        logger.info(f'🎯 Finish reason: {response.finish_reason}')
-        
-    except Exception as e:
-        logger.error(f'❌ Creative writing failed: {e}')
-
-
-async def example_code_generation():
-    """Example: Code generation."""
-    logger.info('🚀 Running code generation example')
-    
-    client = VLLMClient()
-    
-    request = VLLMRequest(
-        prompt=(
-            'Write a Python function that calculates the fibonacci '
-            'sequence up to n terms:\n\n```python\n'
-        ),
-        max_tokens=300,
-        temperature=0.2,  # Lower temperature for code
-        stop=['```', '\n\n\n']
-    )
-    
-    try:
-        response = await client.generate_text(request)
-        
-        logger.success('✅ Code generation completed')
-        logger.info(f'💻 Generated code:\n```python\n{response.text}\n```')
-        
-    except Exception as e:
-        logger.error(f'❌ Code generation failed: {e}')
-
-
-async def main():
-    """Run all examples."""
-    logger.info('🎯 Starting VLLM Client Examples')
-    logger.info('=' * 50)
-    
-    examples = [
-        example_basic_generation,
-        example_batch_generation,
-        example_creative_writing,
-        example_code_generation
-    ]
-    
-    for example in examples:
-        await example()
-        logger.info('-' * 50)
-        await asyncio.sleep(1)  # Brief pause between examples
-        
-    logger.info('🎉 All examples completed!')
-
-
-if __name__ == '__main__':
-    # Configure logger
-    logger.remove()
-    logger.add(
-        lambda msg: print(msg, end=''),
-        format='<green>{time:HH:mm:ss}</green> | '
-               '<level>{level: <8}</level> | '
-               '<cyan>{message}</cyan>',
-        level='INFO'
-    )
-    
-    try:
-        asyncio.run(main())
-    except KeyboardInterrupt:
-        logger.info('\n👋 Goodbye!')
-    except Exception as e:
-        logger.error(f'❌ Script failed: {e}')

llm_utils/scripts/requirements_example.txt

@@ -1,3 +0,0 @@
-aiohttp>=3.8.0
-loguru>=0.6.0
-pydantic>=2.0.0

llm_utils/scripts/serve_script.sh

@@ -1,2 +0,0 @@
-HF_HOME=/home/anhvth5/.cache/huggingface CUDA_VISIBLE_DEVICES=0 /home/anhvth5/miniconda3/envs/unsloth_env/bin/vllm serve ./outputs/8B_selfeval_retranslate/Qwen3-8B_2025_05_30/ls_response_only_r8_a8_sq8192_lr5e_06_bz64_ep1_4/ --port 8140 --tensor-parallel 1 --gpu-memory-utilization 0.9 --dtype auto --max-model-len 8192 --enable-prefix-caching --disable-log-requests --served-model-name selfeval_8b
-Logging to /tmp/vllm_8140.txt

speedy_utils-1.0.15.dist-info/RECORD

@@ -1,40 +0,0 @@
-llm_utils/__init__.py,sha256=0y7y1ia3jyvY3qgOnVp1xCqwMT5UYgw_cSAU_9AjHUo,650
-llm_utils/chat_format/__init__.py,sha256=8dBIUqFJvkgQYedxBtcyxt-4tt8JxAKVap2JlTXmgaM,737
-llm_utils/chat_format/display.py,sha256=a3zWzo47SUf4i-uic-dwf-vxtu6gZWLbnJrszjjZjQ8,9801
-llm_utils/chat_format/transform.py,sha256=328V18FOgRQzljAl9Mh8NF4Tl-N3cZZIPmAwHQspXCY,5461
-llm_utils/chat_format/utils.py,sha256=xTxN4HrLHcRO2PfCTR43nH1M5zCa7v0kTTdzAcGkZg0,1229
-llm_utils/group_messages.py,sha256=wyiZzs7O8yK2lyIakV2x-1CrrWVT12sjnP1vVnmPet4,3606
-llm_utils/lm/__init__.py,sha256=e8eCWlLo39GZjq9CokludZGHYVZ7BnbWZ6GOJoiWGzE,110
-llm_utils/lm/alm.py,sha256=_q8dtvUy_JkYrCV83WwSLD5o5t8e2h9eTRW3DJfWbOM,16850
-llm_utils/lm/chat_html.py,sha256=cbrKjHWIblCat5TRe0X-N7kT7bBtutNYQaIEK-myLq4,8543
-llm_utils/lm/lm.py,sha256=BqACb05k5yV7fU4PLfTuA8e21wXnpseh6OPLKjOlhV0,31185
-llm_utils/lm/lm_json.py,sha256=At_sHZJRU_7_FtvpZ9HbkKLfftO6DkDNhRc9XVoUfx4,2039
-llm_utils/lm/utils.py,sha256=-fDNueiXKQI6RDoNHJYNyORomf2XlCf2doJZ3GEV2Io,4762
-llm_utils/scripts/README.md,sha256=yuOLnLa2od2jp4wVy3rV0rESeiV3o8zol5MNMsZx0DY,999
-llm_utils/scripts/example_vllm_client.py,sha256=tmSRlqHOzBEQ7nTq5I_pY4SH1QKhZF_OMjgD_PV3zp8,8165
-llm_utils/scripts/requirements_example.txt,sha256=mN0Qckft6kUCyuqnvGp1Q3N_RHrbg64AlMvZ-2ougOc,45
-llm_utils/scripts/serve_script.sh,sha256=W9V-wfdWbAWl2riYcXvtYtP5TtmAUss4z8uwZf8wZXs,428
-llm_utils/scripts/vllm_load_balancer.py,sha256=17zaq8RJseikHVoAibGOz0p_MCLcNlnhZDkk7g4cuLc,17519
-llm_utils/scripts/vllm_serve.py,sha256=CbW_3Y9Vt7eQYoGGPT3yj1nhbLYOc3b1LdJBy1sVX-Y,11976
-speedy_utils/__init__.py,sha256=kxQk4PGS3Xkxnerm0YqjF6GKTpgoaTc1vudKid-2c_A,3388
-speedy_utils/all.py,sha256=A9jiKGjo950eg1pscS9x38OWAjKGyusoAN5mrfweY4E,3090
-speedy_utils/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-speedy_utils/common/clock.py,sha256=3n4FkCW0dz46O8By09V5Pve1DSMgpLDRbWEVRryryeQ,7423
-speedy_utils/common/function_decorator.py,sha256=r_r42qCWuNcu0_aH7musf2BWvcJfgZrD81G28mDcolw,2226
-speedy_utils/common/logger.py,sha256=NIOlhhACpcc0BUTSJ8oDYrLp23J2gW_KJFyRVdLN2tY,6432
-speedy_utils/common/notebook_utils.py,sha256=7tFXMaE365ubrD9VCwDoZe7oBqZ2qG5t_Vla4Qvg1M8,2062
-speedy_utils/common/report_manager.py,sha256=dgGfS_fHbZiQMsLzkgnj0OfB758t1x6B4MhjsetSl9A,3930
-speedy_utils/common/utils_cache.py,sha256=yHtoe29Zpzgkp9sGknkIVctm6qcRa6BMR5x6PhOCzvk,8384
-speedy_utils/common/utils_io.py,sha256=d7PKz5tOPrwHDr7GXuYmILvjXJOFEwfzAEIuUcYaI60,4790
-speedy_utils/common/utils_misc.py,sha256=nsQOu2jcplcFHVQ1CnOjEpNcdxIINveGxB493Cqo63U,1812
-speedy_utils/common/utils_print.py,sha256=iQqnOYw2EFC8TqeSDbrcnIQAUKT7FbB8Mec8b2aGAzw,4833
-speedy_utils/multi_worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-speedy_utils/multi_worker/process.py,sha256=f5DKWysQqxSHwEZ8RobxqxfyMPMb712_raynYb2cPmc,6867
-speedy_utils/multi_worker/thread.py,sha256=9pXjvgjD0s0Hp0cZ6I3M0ndp1OlYZ1yvqbs_bcun_Kw,12775
-speedy_utils/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-speedy_utils/scripts/mpython.py,sha256=_J3Y38FsMcCNf-ZWSZ9iOt1Kqed9gr0WrmMHRleaeVY,3871
-speedy_utils/scripts/openapi_client_codegen.py,sha256=4Y1HO8Uht-hCfHEVxY6TlSeTOcZZRVwJ-SBu-bFsAOk,9747
-speedy_utils-1.0.15.dist-info/METADATA,sha256=AeZXiLOUj76H0kAiduMoozmo30SjlFSgus6cg23bG7k,7392
-speedy_utils-1.0.15.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-speedy_utils-1.0.15.dist-info/entry_points.txt,sha256=T1t85jwx8fK6m5msdkBGIXH5R5Kd0zSL0S6erXERPzg,237
-speedy_utils-1.0.15.dist-info/RECORD,,