speedy-utils 1.0.14__py3-none-any.whl → 1.0.16__py3-none-any.whl

llm_utils/lm/lm.py

@@ -4,25 +4,24 @@ import base64
 import hashlib
 import json
 import os
-from token import OP
+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 huggingface_hub import repo_info
 from loguru import logger
-from numpy import isin
-from openai import OpenAI, AuthenticationError, RateLimitError
+from openai import AuthenticationError, OpenAI, RateLimitError
 from openai.pagination import SyncPage
 from openai.types.chat import (
     ChatCompletionAssistantMessageParam,
@@ -31,10 +30,11 @@ 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
-import warnings
+
+from llm_utils.chat_format.display import get_conversation_one_turn
+from speedy_utils.common.utils_io import jdumps
 
 # --------------------------------------------------------------------------- #
 # type helpers
@@ -68,6 +68,17 @@ def _yellow(text: str) -> str:
     return f"\x1b[33m{text}\x1b[0m"
 
 
+# from functools import lru_cache
+
+
+# @lru_cache(maxsize=10)
+# def get_tok(tokenizer_name):
+#     from transformers import AutoTokenizer
+
+#     tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True)
+#     return tokenizer
+
+
 class LM:
     """
     Unified language-model wrapper.
@@ -159,7 +170,7 @@ class LM:
                 available_models = self.list_models(port=port)
                 if available_models:
                     self.model = available_models[0]
-                    logger.info(f"Auto-selected model: {self.model}")
+                    logger.debug(f"Auto-selected model: {self.model}")
                 else:
                     raise ValueError("No models available to select from.")
             else:
@@ -258,9 +269,9 @@ class LM:
             if parsed:
                 # print(_green('<Parsed Structure>'))
                 if hasattr(parsed, "model_dump"):
-                    print(json.dumps(parsed.model_dump(), indent=2))
+                    print(jdumps(parsed.model_dump(), indent=2))
                 else:
-                    print(json.dumps(parsed, indent=2))
+                    print(jdumps(parsed, indent=2))
                 # print(_green('</Parsed Structure>'))
                 print()
 
@@ -468,91 +479,388 @@ class LM:
             return None
 
     @staticmethod
-    def list_models(port=None, host="localhost") -> List[str]:
-        """
-        List available models.
-        """
+    def list_models(
+        port=None, host="localhost", base_url: Optional[str] = None
+    ) -> List[str]:
+        """List available models from OpenAI-compatible API server."""
         try:
-            client: OpenAI = LM(port=port, host=host).client
-            base_url: URL = client.base_url
-            logger.debug(f"Base URL: {base_url}")
+            client: OpenAI = OpenAI(
+                api_key=os.getenv("OPENAI_API_KEY", "abc"),
+                base_url=f"http://{host}:{port}/v1" if port else base_url or None,
+            )
             models: SyncPage[Model] = client.models.list()
             return [model.id for model in models.data]
         except Exception as exc:
-            logger.error(f"Failed to list models: {exc}")
-            return []
+            endpoint = f"http://{host}:{port}/v1" if port else base_url
+            error_msg = str(exc)
 
+            if "404" in error_msg or "Not Found" in error_msg:
+                raise ValueError(
+                    f"No OpenAI-compatible API found at {endpoint}. "
+                    f"The endpoint appears to be running a different service "
+                    f"(possibly Jupyter Server). Please check the port number."
+                ) from exc
+            elif "Connection" in error_msg:
+                raise ValueError(
+                    f"Cannot connect to {endpoint}. "
+                    f"Please verify the service is running and accessible."
+                ) from exc
+            else:
+                raise ValueError(
+                    f"Failed to list models from {endpoint}: {error_msg}"
+                ) from exc
 
-from functools import cache
-from llm_utils.lm.lm import LM, RawMsgs
-from pydantic import BaseModel
-import re
-import json
-from typing import *
-import re
+    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,
+    ):
+        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 = self.client.chat.completions.create(
+            model=self.model,  # type: ignore
+            messages=messages,  # type: ignore
+            extra_body={"guided_json": json_schema},
+            **model_kwargs,
+        )
 
-class LMReasoner(LM):
-    "Regex-based reasoning wrapper for LM."
+        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 build_regex_from_pydantic(self, model: type[BaseModel]) -> str:
+    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
+
+    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]:
         """
-        Build a regex pattern string for validating output that should match a Pydantic model.
+        Inspect word probabilities in a language model response.
 
         Args:
-            model: A Pydantic BaseModel class
+            tokenizer: Tokenizer instance to encode words.
+            messages: List of messages to analyze.
 
         Returns:
-            A regex string that matches a JSON representation of the model
+            A tuple containing:
+            - List of word probabilities with their log probabilities.
+            - Token log probability dictionaries.
+            - Rendered string with colored word probabilities.
         """
-        # regex = f"<think>\\n.*?\\n</think>\\n\\n\\```json\\n.*"
-        print(f"{regex=}")
+        if messages is None:
+            messages = 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 = inspect_word_probs(self, tokenizer, messages)
+        if do_print:
+            print(ret[-1])
+        return ret
+
+    def last_messages(self, add_think: bool = True) -> Optional[List[Dict[str, str]]]:
+        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
 
-        return regex
 
-    def __call__(
-        self,
-        response_format: type[BaseModel],
-        prompt: Optional[str] = None,
-        messages: Optional[RawMsgs] = None,
-        **kwargs,
-    ):
+@lru_cache(maxsize=10)
+def get_tokenizer(model_name: str) -> Any:
+    from transformers import AutoTokenizer # type: ignore
 
-        if prompt is not None:
-            output = super().__call__(
-                prompt=prompt
-                + "\nresponse_format:"
-                + str(response_format.model_json_schema()),
-                response_format=str,
-                # extra_body={"guided_regex": regex},
-                **kwargs,
-            )  # type: ignore
-        elif messages is not None:
-            # append last message with the json schema
-            messages[-1]["content"] += "\nresponse_format:" + str(  # type: ignore
-                response_format.model_json_schema()
-            )
-            output = super().__call__(
-                messages=messages,
-                response_format=str,
-                # extra_body={"guided_regex": regex},
-                **kwargs,
-            )
-        else:
-            raise ValueError("Either prompt or messages must be provided.")
-        # import ipdb; ipdb.set_trace()
-        # parse using regex
-        pattern = re.compile(
-            r"<think>\n(?P<think>.*?)\n</think>\n\n(?P<json>\{.*\})",
-            re.DOTALL,
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    return tokenizer
+
+
+def inspect_word_probs(lm, tokenizer, messages):
+
+    import numpy as np
+
+    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 = lm_client.client.completions.create(
+            model=lm_client.model,  # type: ignore
+            prompt=prompt,
+            max_tokens=1,
+            logprobs=1,
+            extra_body={"prompt_logprobs": 0},
         )
-        match = pattern.search(output)
-        if not match:
-            raise ValueError("Output does not match expected format")
-        parsed_output = match.group(0)
-        think_part = match.group("think")
-        json_part = match.group("json")
+        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
 
-        pydantic_object = response_format.model_validate(json.loads(json_part))
-        return pydantic_object
+        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 = compute_word_log_probs(tokenizer, lm)
+    return word_probs, token_logprob_dicts, render_by_logprob(word_probs)
+
+
+class LLMTask(ABC):
+    """
+    Callable wrapper around an LM endpoint.
+
+    Sub-classes must set:
+      • lm              – the 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(LLMTask):
+            "TODO: SYSTEM_PROMPT_INSTURCTION HERE"
+
+            lm = LM(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()
+        demo_task({'text_to_translate': 'Translate from english to vietnamese: Hello how are you'})
+    ```
+    """
+
+    lm: "LM"
+    InputModel: Type[BaseModel]
+    OutputModel: Type[BaseModel]
+
+    temperature: float = 0.6
+    think: bool = False
+    add_json_schema: bool = False
+
+    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 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]
+        return get_conversation_one_turn(
+            system_msg=system_prompt, user_msg=user_msg, assistant_msg=assistant_msg
+        )

llm_utils/lm/lm_json.py

@@ -0,0 +1,68 @@
+from typing import Any, Optional
+
+
+from llm_utils.lm.lm import LM, RawMsgs
+
+
+class LMJson(LM):
+    "Regex-based reasoning wrapper for LM."
+
+    def __init__(
+        self,
+        model: str | None = None,
+        *,
+        temperature: float = 0.0,
+        max_tokens: int = 2_000,
+        host: str = "localhost",
+        port: Optional[int | str] = None,
+        base_url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        cache: bool = True,
+        **openai_kwargs: Any,
+    ) -> None:
+        """
+        Initialize the LMJson instance.
+
+        Args:
+            model (str | None): The model name to use.
+            temperature (float): Sampling temperature.
+            max_tokens (int): Maximum number of tokens to generate.
+            host (str): Host for the API.
+            port (int | str, optional): Port for the API.
+            base_url (str, optional): Base URL for the API.
+            api_key (str, optional): API key for authentication.
+            cache (bool): Whether to cache responses.
+            **openai_kwargs: Additional OpenAI parameters.
+        """
+        super().__init__(
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            host=host,
+            port=port,
+            base_url=base_url,
+            api_key=api_key,
+            cache=cache,
+            **openai_kwargs,
+        )
+
+    def __call__(
+        self,
+        prompt: Optional[str] = None,
+        messages: Optional[RawMsgs] = None,
+        cache: Optional[bool] = None,
+        max_tokens: Optional[int] = None,
+        return_openai_response: bool = False,
+        **kwargs: Any,
+    ):
+
+        output = super().__call__(
+            prompt=prompt,
+            messages=messages,
+            response_format=str,
+            cache=cache,
+            max_tokens=max_tokens,
+            return_openai_response=return_openai_response,
+            **kwargs,
+        )
+        return output

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/README.md

@@ -0,0 +1,48 @@
+# VLLM Server Examples
+
+This directory contains scripts for working with VLLM servers.
+
+## Files
+
+- `serve_script.sh` - Script to start the VLLM server
+- `example_vllm_client.py` - Beautiful example client for interacting with VLLM
+- `requirements_example.txt` - Python dependencies for the example
+
+## Usage
+
+### 1. Start the VLLM Server
+
+```bash
+bash serve_script.sh
+```
+
+### 2. Install Dependencies
+
+```bash
+pip install -r requirements_example.txt
+```
+
+### 3. Run Examples
+
+```bash
+python example_vllm_client.py
+```
+
+## Features
+
+The example client demonstrates:
+
+- ✅ Basic text generation
+- ✅ Batch processing
+- ✅ Creative writing with high temperature
+- ✅ Code generation with low temperature
+- ✅ Proper error handling
+- ✅ Health checks
+- ✅ Beautiful logging with loguru
+- ✅ Type safety with Pydantic models
+- ✅ Async/await patterns
+
+## Configuration
+
+The client connects to `http://localhost:8140` by default.
+Modify the `VLLMClient` initialization to use different servers.

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