speedy-utils 1.1.4__py3-none-any.whl → 1.1.6__py3-none-any.whl

llm_utils/chat_format/display.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
+
+from difflib import SequenceMatcher
 from typing import Any, Optional
+
 from IPython.display import HTML, display
-from difflib import SequenceMatcher
 
 
 def show_chat(
@@ -19,6 +21,17 @@ def show_chat(
         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."
 
+    if isinstance(msgs[-1], dict) and "choices" in msgs[-1]:
+        message = msgs[-1]["choices"][0]["message"]
+        reasoning_content = message.get("reasoning_content")
+        content = message.get("content", "")
+        if reasoning_content:
+            content = reasoning_content + "\n" + content
+        msgs[-1] = {
+            "role": message["role"],
+            "content": content,
+        }
+
     themes: dict[str, dict[str, dict[str, str]]] = {
         "default": {
             "system": {"background": "#ffaaaa", "text": "#222222"},  # More red
@@ -156,9 +169,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
 

llm_utils/lm/async_lm/__init__.py

@@ -0,0 +1,2 @@
+from .async_lm import AsyncLM
+from .async_llm_task import AsyncLLMTask

llm_utils/lm/async_lm/_utils.py

@@ -0,0 +1,198 @@
+from functools import lru_cache
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    List,
+    TypeVar,
+    Union,
+)
+
+# from openai.pagination import AsyncSyncPage
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+)
+from pydantic import BaseModel
+from typing_extensions import TypedDict
+
+# --------------------------------------------------------------------------- #
+# type helpers
+# --------------------------------------------------------------------------- #
+TModel = TypeVar("TModel", bound=BaseModel)
+Messages = List[ChatCompletionMessageParam]
+LegacyMsgs = List[Dict[str, str]]
+RawMsgs = Union[Messages, LegacyMsgs]
+
+# --------------------------------------------------------------------------- #
+# color helpers (unchanged)
+# --------------------------------------------------------------------------- #
+
+
+def _color(code: int, text: str) -> str:
+    return f"\x1b[{code}m{text}\x1b[0m"
+
+
+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)
+
+
+TParsed = TypeVar("TParsed", bound=BaseModel)
+
+
+class ParsedOutput(TypedDict, Generic[TParsed]):
+    messages: List
+    completion: Any
+    parsed: TParsed
+
+
+# --------------------------------------------------------------------------- #
+# 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)
+
+
+__all__ = [
+    "TModel",
+    "Messages",
+    "LegacyMsgs",
+    "RawMsgs",
+    "TParsed",
+    "ParsedOutput",
+    "get_tokenizer",
+    "inspect_word_probs_async",
+    "_color",
+    "_red",
+    "_green",
+    "_blue",
+    "_yellow",
+]
+# --------------------------------------------------------------------------- #]

llm_utils/lm/async_lm/async_llm_task.py

@@ -0,0 +1,154 @@
+from abc import ABC
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    TypeVar,
+    Union,
+    cast,
+)
+
+# from openai.pagination import AsyncSyncPage
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+)
+from pydantic import BaseModel
+
+from llm_utils.chat_format.display import get_conversation_one_turn
+
+from .async_lm import AsyncLM
+
+# --------------------------------------------------------------------------- #
+# type helpers
+# --------------------------------------------------------------------------- #
+TModel = TypeVar("TModel", bound=BaseModel)
+Messages = List[ChatCompletionMessageParam]
+LegacyMsgs = List[Dict[str, str]]
+RawMsgs = Union[Messages, LegacyMsgs]
+
+
+# --------------------------------------------------------------------------- #
+# Async LLMTask class
+# --------------------------------------------------------------------------- #
+
+InputModelType = TypeVar("InputModelType", bound=BaseModel)
+OutputModelType = TypeVar("OutputModelType", bound=BaseModel)
+
+
+class AsyncLLMTask(ABC, Generic[InputModelType, OutputModelType]):
+    """
+    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: InputModelType
+    OutputModel: OutputModelType
+
+    temperature: float = 0.6
+    think: bool = False
+    add_json_schema: bool = False
+    cache: bool = False
+
+    async def __call__(
+        self,
+        data: BaseModel | dict,
+        temperature: float = 0.1,
+        cache: bool = False,
+        think: Optional[bool] = None,  # if not None, overrides self.think
+    ) -> tuple[OutputModelType, List[Dict[str, Any]]]:
+        # Get the input and output model types from the generic parameters
+        type_args = getattr(self.__class__, "__orig_bases__", None)
+        if (
+            type_args
+            and hasattr(type_args[0], "__args__")
+            and len(type_args[0].__args__) >= 2
+        ):
+            input_model = type_args[0].__args__[0]
+            output_model = type_args[0].__args__[1]
+        else:
+            # Fallback to the old way if type introspection fails
+            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 or use proper generic typing."
+                )
+            input_model = self.InputModel
+            output_model = self.OutputModel
+
+        # Ensure input_model is a class before calling
+        if isinstance(data, BaseModel):
+            item = data
+        elif isinstance(input_model, type) and issubclass(input_model, BaseModel):
+            item = input_model(**data)
+        else:
+            raise TypeError("InputModel must be a subclass of BaseModel")
+
+        assert isinstance(output_model, type) and issubclass(output_model, BaseModel), (
+            "OutputModel must be a subclass of BaseModel"
+        )
+
+        result = await self.lm.parse(
+            prompt=item.model_dump_json(),
+            instruction=self.__doc__ or "",
+            response_model=output_model,
+            temperature=temperature or self.temperature,
+            think=think if think is not None else self.think,
+            add_json_schema_to_instruction=self.add_json_schema,
+            cache=self.cache or cache,
+        )
+
+        return (
+            cast(OutputModelType, result["parsed"]),  # type: ignore
+            cast(List[dict], result["messages"]),  # type: ignore
+        )
+
+    def generate_training_data(
+        self, input_dict: Dict[str, Any], output: Dict[str, Any]
+    ) -> 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}
+
+    arun = __call__  # alias for compatibility with other LLMTask implementations