speedy-utils 2.0.1__tar.gz → 2.0.2__tar.gz

{speedy_utils-2.0.1 → speedy_utils-2.0.2}/CLAUDE.md

@@ -9,7 +9,7 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 uv sync
 
 # Run tests
-uv run pytest
+uv run pytest -n 32
 
 # Run a single test file
 uv run pytest tests/test_thread.py
@@ -60,21 +60,41 @@ def some_function():
 
 | Module | Import Time | Strategy |
 |--------|-------------|----------|
-| `torch` | ~5s | Lazy import in GPU functions |
-| `matplotlib` | ~3.6s | Lazy import in plotting functions |
-| `pandas` | ~1.3s | Lazy import in data functions |
-| `IPython` | ~1s | Lazy import in notebook functions |
+| `torch` | ~5s | Lazy import inside functions |
+| `matplotlib` | ~1.3s cumulative | Lazy via `__getattr__` in `__init__.py` |
+| `pandas` | ~1.9s cumulative | Lazy via `__getattr__` in `__init__.py` |
+| `IPython` | ~2.1s cumulative | Lazy via `__getattr__` in `__init__.py` |
 
-### Using the _LazyModule Pattern
+### Rule: Never top-level import heavy modules in `__init__.py`
+
+Only a short `__getattr__` block (keyed on module name) is allowed for heavy
+modules. All internal `speedy_utils.*` imports must remain direct top-level
+imports. External heavy modules (`pandas`, `matplotlib`, `IPython`, `torch`)
+must **never** appear as top-level imports in any `__init__.py`.
 
 ```python
-from speedy_utils.__imports import pd, plt
+# BAD — in __init__.py
+import pandas as pd          # adds ~1.3 s
+
+# GOOD — in __init__.py
+# pd is declared in _HEAVY dict, resolved by __getattr__ on first access
 
-def process_data():
-    df = pd.DataFrame()  # pandas imports here, not at module load
-    return df
+# GOOD — inside a function anywhere in the package
+def load_dataframe(path):
+    import pandas as pd      # only pays cost when called
+    return pd.read_csv(path)
 ```
 
+### Rule: Heavy lazy names must NOT be in `__all__`
+
+`from speedy_utils import *` resolves every name in `__all__`, which triggers
+`__getattr__` for each heavy name and defeats the lazy-load strategy.  Only
+include names that are directly importable at module load time in `__all__`.
+Heavy names (`np`, `pd`, `matplotlib`, `plt`, `get_ipython`, `HTML`, `display`,
+`BaseModel`) must be kept out of `__all__`; they remain accessible via explicit
+attribute access (`speedy_utils.pd`) or explicit import
+(`from speedy_utils import pd`).
+
 ## Architecture
 
 ### Package Structure

{speedy_utils-2.0.1 → speedy_utils-2.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: speedy-utils
-Version: 2.0.1
+Version: 2.0.2
 Summary: Fast and easy-to-use package for data science
 Project-URL: Homepage, https://github.com/anhvth/speedy
 Project-URL: Repository, https://github.com/anhvth/speedy
@@ -24,6 +24,7 @@ Requires-Dist: cachetools
 Requires-Dist: chainlit>=2.0
 Requires-Dist: datasets
 Requires-Dist: debugpy
+Requires-Dist: dill
 Requires-Dist: fastcore
 Requires-Dist: fastprogress
 Requires-Dist: freezegun
@@ -648,6 +649,36 @@ answer = message.content
 reasoning = getattr(message, "reasoning_content", None)
 ```
 
+For custom staged generation with explicit tags, `Qwen3LLM` also exposes
+`complete_until()`:
+
+```python
+from llm_utils import Qwen3LLM
+
+llm = Qwen3LLM(model="Qwen/Qwen3-0.6B")
+
+memory_state = llm.complete_until(
+    [{"role": "user", "content": "Solve this in stages"}],
+    "<memory>",
+    stop="</memory>",
+    max_tokens=256,
+)
+
+think_state = llm.complete_until(
+    [{"role": "user", "content": "Solve this in stages"}],
+    memory_state.assistant_prompt_prefix + "\n<think_efficient>",
+    stop="</think_efficient>",
+    max_tokens=512,
+)
+
+final_state = llm.complete_until(
+    [{"role": "user", "content": "Solve this in stages"}],
+    think_state.assistant_prompt_prefix,
+    stop="<|im_end|>",
+    max_tokens=512,
+)
+```
+
 #### Conversation History
 
 Inspect previous conversations:

{speedy_utils-2.0.1 → speedy_utils-2.0.2}/README.md

@@ -597,6 +597,36 @@ answer = message.content
 reasoning = getattr(message, "reasoning_content", None)
 ```
 
+For custom staged generation with explicit tags, `Qwen3LLM` also exposes
+`complete_until()`:
+
+```python
+from llm_utils import Qwen3LLM
+
+llm = Qwen3LLM(model="Qwen/Qwen3-0.6B")
+
+memory_state = llm.complete_until(
+    [{"role": "user", "content": "Solve this in stages"}],
+    "<memory>",
+    stop="</memory>",
+    max_tokens=256,
+)
+
+think_state = llm.complete_until(
+    [{"role": "user", "content": "Solve this in stages"}],
+    memory_state.assistant_prompt_prefix + "\n<think_efficient>",
+    stop="</think_efficient>",
+    max_tokens=512,
+)
+
+final_state = llm.complete_until(
+    [{"role": "user", "content": "Solve this in stages"}],
+    think_state.assistant_prompt_prefix,
+    stop="<|im_end|>",
+    max_tokens=512,
+)
+```
+
 #### Conversation History
 
 Inspect previous conversations:

{speedy_utils-2.0.1 → speedy_utils-2.0.2}/debug/integration_test.py

@@ -23,8 +23,8 @@ print()
 print('✓ Test 2: Generate from text')
 result = lm.generate('The answer is', max_tokens=10, temperature=0.5)
 print(f'  Input: "The answer is"')
-print(f'  Output: {result["text"][:50]}')
-assert 'text' in result
+print(f'  Output: {result.text[:50]}')
+assert result.text is not None
 print()
 
 # Test 3: Generate from token IDs
@@ -32,10 +32,10 @@ print('✓ Test 3: Generate from token IDs')
 input_ids = lm.encode('Python is')
 result = lm.generate(input_ids, max_tokens=15, return_token_ids=True)
 print(f'  Input IDs: {input_ids}')
-print(f'  Output: {result["text"][:50]}')
-print(f'  Output IDs: {result.get("token_ids", "N/A")}')
-assert 'text' in result
-assert 'token_ids' in result
+print(f'  Output: {result.text[:50]}')
+print(f'  Output IDs: {getattr(result, "token_ids", "N/A")}')
+assert result.text is not None
+assert hasattr(result, "token_ids") and result.token_ids is not None
 print()
 
 # Test 4: Multiple generations
@@ -43,7 +43,7 @@ print('✓ Test 4: Multiple generations (n=3)')
 results = lm.generate('AI is', max_tokens=10, n=3, temperature=0.8)
 print(f'  Input: "AI is"')
 for i, r in enumerate(results, 1):
-    print(f'  {i}. {r["text"][:40]}')
+    print(f'  {i}. {r.text[:40]}')
 # Note: Server may return fewer choices than requested
 assert isinstance(results, list)
 assert len(results) >= 1
@@ -53,8 +53,8 @@ print()
 print('✓ Test 5: Temperature control')
 low_temp = lm.generate('1 + 1 =', max_tokens=5, temperature=0.1, seed=42)
 high_temp = lm.generate('1 + 1 =', max_tokens=5, temperature=1.5, seed=123)
-print(f'  Low temp (0.1): {low_temp["text"][:30]}')
-print(f'  High temp (1.5): {high_temp["text"][:30]}')
+print(f'  Low temp (0.1): {low_temp.text[:30]}')
+print(f'  High temp (1.5): {high_temp.text[:30]}')
 print()
 
 # Test 6: Stop sequences
@@ -66,16 +66,16 @@ result = lm.generate(
     temperature=0.7,
 )
 print(f'  Input: "Count: 1, 2, 3,"')
-print(f'  Output (stop at ", 6"): {result["text"][:50]}')
+print(f'  Output (stop at ", 6"): {result.text[:50]}')
 print()
 
 # Test 7: Seed reproducibility
 print('✓ Test 7: Seed reproducibility')
 r1 = lm.generate('Test', max_tokens=10, seed=999, temperature=0.8)
 r2 = lm.generate('Test', max_tokens=10, seed=999, temperature=0.8)
-print(f'  Run 1: {r1["text"][:40]}')
-print(f'  Run 2: {r2["text"][:40]}')
-print(f'  Match: {r1["text"] == r2["text"]}')
+print(f'  Run 1: {r1.text[:40]}')
+print(f'  Run 2: {r2.text[:40]}')
+print(f'  Match: {r1.text == r2.text}')
 print()
 
 print('=' * 60)

{speedy_utils-2.0.1 → speedy_utils-2.0.2}/docs/GENERATE_QUICKREF.md

@@ -2,7 +2,7 @@
 
 ## Summary
 
-The `LLM.generate()` method provides a HuggingFace Transformers-style interface for low-level text generation, working directly with token IDs.
+The `LLM.generate()` method provides a HuggingFace Transformers-style interface for prompt continuation via the completions API. For chat assistant-turn generation, use `LLM.chat_completion()`.
 
 ## Basic Signature
 
@@ -16,9 +16,17 @@ lm.generate(
 ) -> dict | list[dict]
 ```
 
+## Which method to use
+
+- `generate()`: Continue a raw prompt through the completions API.
+- `chat_completion()`: Generate the next assistant turn from chat messages.
+- `pydantic_parse()`: Parse structured output from the chat completions API.
+
+`enable_thinking` is primarily a chat-path control. `generate()` forwards it only for backends that support the same flag on the completions endpoint.
+
 ## Common Use Cases
 
-### 1. Simple Generation
+### 1. Simple Prompt Continuation
 ```python
 result = lm.generate('Hello world', max_tokens=50)
 print(result['text'])

speedy_utils-2.0.2/docs/MIGRATION_THINKING.md

@@ -0,0 +1,73 @@
+# Qwen3LLM Thinking Guide
+
+`Qwen3LLM` is the only model-specific thinking helper kept in this repository.
+It supports prefix continuation for Qwen3-style reasoning traces.
+It also includes a raw prefix-conditioned helper for custom staged tag flows.
+
+## Usage
+
+```python
+from llm_utils import Qwen3LLM
+
+llm = Qwen3LLM(client=8000)
+
+message = llm.chat_completion(
+    [{"role": "user", "content": "Solve x^2 + 2x + 1 = 0"}],
+    thinking_max_tokens=32,
+    content_max_tokens=128,
+)
+print(message)
+```
+
+## Custom Prefix Flows
+
+Use `complete_until()` when you want to generate a custom tagged block and stop
+at a specific closing tag instead of the built-in `<think>...</think>` flow.
+
+```python
+from llm_utils import Qwen3LLM
+
+llm = Qwen3LLM(client=8000)
+
+memory_state = llm.complete_until(
+    [{"role": "user", "content": "Plan the answer in stages"}],
+    "<memory>",
+    stop="</memory>",
+    max_tokens=128,
+)
+
+think_state = llm.complete_until(
+    [{"role": "user", "content": "Plan the answer in stages"}],
+    memory_state.assistant_prompt_prefix + "\n<think_efficient>",
+    stop="</think_efficient>",
+    max_tokens=256,
+)
+
+final_state = llm.complete_until(
+    [{"role": "user", "content": "Plan the answer in stages"}],
+    think_state.assistant_prompt_prefix,
+    stop="<|im_end|>",
+    max_tokens=256,
+)
+```
+
+## Behavior
+
+- `chat_completion()` continues from a partial assistant prefix.
+- The default prefix starts a `<think>` block.
+- `thinking_max_tokens` caps the reasoning phase.
+- `content_max_tokens` caps the final answer phase.
+- `complete_until()` continues from a raw assistant prefix and stops on custom
+  tag boundaries.
+- The returned `assistant_prompt_prefix` includes the generated text and the
+  matched stop token when generation ends on a stop sequence.
+
+## Notes
+
+- Use `n=1`; prefix continuation is single-path only.
+- The returned object is an OpenAI `ChatCompletionMessage`.
+- The intermediate prefix state is internal; the returned value is an OpenAI
+  `ChatCompletionMessage`.
+- The returned message includes `call_count`, which is `1` when the model
+  finishes in one staged call and `2` when it needs a reasoning call plus a
+  follow-up content call.

{speedy_utils-2.0.1 → speedy_utils-2.0.2}/docs/TOKENIZATION.md

@@ -1,12 +1,16 @@
 # Tokenization and Generation Support in LLM
 
-The `LLM` class now includes built-in tokenization and low-level generation support through the `TokenizationMixin`, providing methods similar to HuggingFace Transformers.
+The `LLM` class includes built-in tokenization plus two distinct generation styles:
+- `generate()` for prompt continuation via the completions API
+- `chat_completion()` for assistant-turn generation via the chat completions API
 
 ## Features
 
 - **encode()**: Convert text to token IDs
 - **decode()**: Convert token IDs back to text
-- **generate()**: HuggingFace-style generation with token-level control
+- **generate()**: HuggingFace-style prompt continuation with token-level control
+- **chat_completion()**: Chat-based assistant-turn generation from messages
+- **pydantic_parse()**: Structured chat output parsing
 - Support for special tokens (BOS, EOS, etc.)
 - Optional token string output for debugging
 
@@ -94,6 +98,14 @@ def generate(
 - `dict`: Generation result with 'text', 'token_ids', 'finish_reason' (if n=1)
 - `list[dict]`: List of generation results (if n>1)
 
+### Which method to use
+
+- Use `generate()` when you already have a raw prompt string and want the completions API to continue it with more tokens.
+- Use `chat_completion()` when you want the chat completions API to produce the next assistant turn from messages.
+- Use `pydantic_parse()` when you want structured output from the chat completions API.
+
+`enable_thinking` is mainly meaningful on the chat-based paths. `generate()` passes it through for backends that honor `extra_body.chat_template_kwargs.enable_thinking` on the completions endpoint, but that support is backend-specific.
+
 ## Usage Examples
 
 ### Basic Encoding/Decoding
@@ -169,9 +181,9 @@ result = lm.decode(combined)
 print(result)  # 'Helloworld'
 ```
 
-## HuggingFace-Style Generation
+## Prompt Continuation with `generate()`
 
-The `generate()` method provides a low-level interface similar to HuggingFace Transformers:
+The `generate()` method provides a low-level prompt-continuation interface similar to HuggingFace Transformers. It uses the completions API rather than the chat completions API.
 
 ### Basic Text Generation
 
@@ -184,6 +196,21 @@ result = lm.generate(
 print(result['text'])  # ' Paris.'
 ```
 
+## Chat Assistant-Turn Generation
+
+Use `chat_completion()` when you want the next assistant message from a list of chat messages:
+
+```python
+message = lm.chat_completion(
+    [
+        {"role": "system", "content": "Be concise."},
+        {"role": "user", "content": "What is the capital of France?"},
+    ],
+    temperature=0.2,
+)
+print(message.content)  # 'Paris.'
+```
+
 ### Working with Token IDs
 
 ```python

{speedy_utils-2.0.1 → speedy_utils-2.0.2}/examples/generate_example.py

@@ -1,7 +1,7 @@
-"""Example: Using the HuggingFace-style generate() method.
+"""Example: Using the prompt-continuation `generate()` method.
 
-This demonstrates the low-level generation interface that works
-directly with token IDs, similar to HuggingFace Transformers.
+This demonstrates the low-level completions-style generation interface
+for continuing a raw prompt, similar to HuggingFace Transformers.
 """
 
 from llm_utils.lm import LLM
@@ -10,7 +10,7 @@ from llm_utils.lm import LLM
 lm = LLM(client=8000)
 
 print('=' * 70)
-print('HuggingFace-Style Generation Examples')
+print('Prompt Continuation Examples')
 print('=' * 70)
 print()
 
@@ -23,7 +23,7 @@ result = lm.generate(
     temperature=0.8,
 )
 print(f"Prompt: 'Write a haiku about coding:'")
-print(f"Generated:\n{result['text']}")
+print(f"Generated:\n{result.text}")
 print()
 
 # Example 2: Working with token IDs
@@ -40,8 +40,8 @@ result = lm.generate(
     temperature=0.7,
     return_token_ids=True,  # Get token IDs back
 )
-print(f'Output IDs: {result["token_ids"]}')
-print(f'Output text: {result["text"]}')
+print(f'Output IDs: {result.token_ids}')
+print(f'Output text: {result.text}')
 print()
 
 # Example 3: Temperature sampling
@@ -52,12 +52,12 @@ prompt = 'The best way to learn programming is'
 # Low temperature (more deterministic)
 result_low = lm.generate(prompt, max_tokens=40, temperature=0.1)
 print(f'Temperature 0.1 (deterministic):')
-print(f'  {result_low["text"][:100]}...')
+print(f'  {result_low.text[:100]}...')
 
 # High temperature (more creative)
 result_high = lm.generate(prompt, max_tokens=40, temperature=1.5)
 print(f'Temperature 1.5 (creative):')
-print(f'  {result_high["text"][:100]}...')
+print(f'  {result_high.text[:100]}...')
 print()
 
 # Example 4: Top-k and Top-p sampling
@@ -71,7 +71,7 @@ result = lm.generate(
     top_p=0.95,  # Nucleus sampling
     repetition_penalty=1.2,  # Reduce repetition
 )
-print(f'Generated: {result["text"]}')
+print(f'Generated: {result.text}')
 print()
 
 # Example 5: Multiple generations (like num_return_sequences)
@@ -84,7 +84,7 @@ results = lm.generate(
     n=4,  # Generate 4 different completions
 )
 for i, result in enumerate(results, 1):
-    print(f'{i}. {result["text"]}')
+    print(f'{i}. {result.text}')
 print()
 
 # Example 6: Controlled generation with stop sequences
@@ -96,7 +96,7 @@ result = lm.generate(
     temperature=0.7,
     stop=['\n\n', 'Instructions:'],  # Stop at these sequences
 )
-print(f'Generated:\n{result["text"]}')
+print(f'Generated:\n{result.text}')
 print()
 
 # Example 7: Reproducible generation with seed
@@ -109,10 +109,10 @@ result1 = lm.generate(prompt, max_tokens=10, temperature=0.8, seed=42)
 result2 = lm.generate(prompt, max_tokens=10, temperature=0.8, seed=42)
 result3 = lm.generate(prompt, max_tokens=10, temperature=0.8, seed=123)
 
-print(f'Seed 42 (run 1): {result1["text"]}')
-print(f'Seed 42 (run 2): {result2["text"]}')
-print(f'Seed 123:        {result3["text"]}')
-print(f'Results match:   {result1["text"] == result2["text"]}')
+print(f'Seed 42 (run 1): {result1.text}')
+print(f'Seed 42 (run 2): {result2.text}')
+print(f'Seed 123:        {result3.text}')
+print(f'Results match:   {result1.text == result2.text}')
 print()
 
 # Example 8: Token count estimation
@@ -132,7 +132,7 @@ result = lm.generate(
     max_tokens=min(100, available_tokens),
     temperature=0.7,
 )
-generated_tokens = lm.encode(result['text'])
+generated_tokens = lm.encode(result.text)
 print(f'Generated tokens: {len(generated_tokens)}')
 print(f'Total tokens used: {len(prompt_tokens) + len(generated_tokens)}')
 print()

{speedy_utils-2.0.1 → speedy_utils-2.0.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "speedy-utils"
-version = "2.0.1"
+version = "2.0.2"
 description = "Fast and easy-to-use package for data science"
 authors = [{ name = "AnhVTH", email = "anhvth.226@gmail.com" }]
 readme = "README.md"
@@ -23,6 +23,7 @@ dependencies = [
     "pydantic",
     "tqdm",
     "cachetools",
+    "dill",
     "ruff",
     "bump2version",
     "json-repair",
@@ -143,3 +144,8 @@ python_classes = ["Test*"]
 python_functions = ["test_*"]
 # Exclude debug and example directories from test discovery
 norecursedirs = ["debug", "examples", "experiments", "notebooks", "scripts", ".git", "__pycache__"]
+
+[dependency-groups]
+dev = [
+    "pytest-xdist>=3.8.0",
+]

speedy_utils-2.0.2/pyrightconfig.json

@@ -0,0 +1,8 @@
+{
+  "include": ["src", "tests"],
+  "exclude": ["debug", "examples", "scripts", "**/node_modules", "**/__pycache__", "**/.*", ".venv"],
+  "diagnosticSeverityOverrides": {
+    "reportUnsupportedDunderAll": "none"
+  },
+  "pythonVersion": "3.11"
+}

{speedy_utils-2.0.1 → speedy_utils-2.0.2}/scripts/bug_simple.py

@@ -4,7 +4,7 @@ from speedy_utils import *
 def do_something(x):
     x = 10
     y = 0
-    x/y
+    x/y  # type: ignore[operator]
 
 
 do_something(1)

{speedy_utils-2.0.1 → speedy_utils-2.0.2}/scripts/test.py

@@ -10,7 +10,7 @@ def demo_1_simple_error():
     def process_item(x):
         # Intentional error - calling a list
         my_list = [1, 2, 3]
-        return my_list(x)
+        return my_list(x)  # type: ignore[operator]
 
     # try:
     multi_thread(process_item, range(300), workers=100, progress=True)

speedy_utils-2.0.2/scripts/test_locals.py

@@ -0,0 +1,17 @@
+from speedy_utils import multi_process
+
+
+def process_data(item):
+    # Mix of user variables and imported modules
+    data = {"value": item}
+    multiplier = 2
+
+    # This will cause an error
+    denominator = 0
+    final = data["value"] * multiplier / denominator
+    return final
+
+
+if __name__ == "__main__":
+    inputs = range(5)
+    multi_process(process_data, inputs, backend="mp")