dustylm-sdk 0.1.1__tar.gz

dustylm_sdk-0.1.1/PKG-INFO

@@ -0,0 +1,143 @@
+Metadata-Version: 2.4
+Name: dustylm-sdk
+Version: 0.1.1
+Summary: The official Python SDK for DustyLM: an 8M-parameter model that talks like a robot vacuum.
+Author-email: Mahmood Khordoo <m.khordoo@gmail.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/khordoo/dustylm-sdk
+Project-URL: Source, https://github.com/khordoo/dustylm-sdk
+Keywords: llm,slm,small-language-model,transformer,pytorch,onnx,inference
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0
+Requires-Dist: tokenizers>=0.15
+Requires-Dist: huggingface-hub>=0.20
+Requires-Dist: numpy>=1.24
+Provides-Extra: onnx
+Requires-Dist: onnxruntime>=1.15; extra == "onnx"
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/khordoo/dustylm-sdk/main/assets/logo.png" alt="DustyLM Logo" width="200"/>
+</p>
+
+# dustylm
+
+Run DustyLM, the 8M parameter robot vacuum language model, in a few lines.
+
+```
+pip install dustylm-sdk
+```
+
+```python
+from dustylm import DustyLM
+model = DustyLM.from_pretrained("mkhordoo/dusty-8m-sft")
+response = model.chat([{"role": "user", "content": "who are you?"}])
+print(response["choices"][0]["message"]["content"])
+
+# beep. i am a little robot. i clean floors and find crumbs.
+```
+
+## Backends
+
+| Backend | Install | Model file | Best for |
+|---|---|---:|---|
+| `torch` (default) | `pip install dustylm-sdk` | ~32 MB (FP32) | Inspection and experimentation |
+| `onnx` | `pip install dustylm-sdk[onnx]` | ~8–10 MB (int8) | A compact model artifact |
+
+These sizes describe the model files, not the complete Python environment.
+Installation size varies by platform. Both backends currently include PyTorch as
+an SDK dependency.
+
+```python
+# ONNX backend
+model = DustyLM.from_pretrained("mkhordoo/dusty-8m-sft", backend="onnx")
+```
+
+## Loading Local Checkpoints
+
+Pass a Hugging Face repository ID to download its model and tokenizer
+automatically, as shown in the quick-start example:
+
+```python
+model = DustyLM.from_pretrained("mkhordoo/dusty-8m-sft")
+```
+
+For files already downloaded to your machine, pass their local directory
+instead. With the default filenames, that directory must contain
+`tokenizer.json` and either `model.pt` (PyTorch) or `model_int8.onnx` (ONNX):
+
+```python
+model = DustyLM.from_pretrained("./my-checkpoint/")
+```
+
+You can use any local directory as the artifact root. `model_file` and
+`tokenizer_file` are resolved relative to that directory, so they may use
+different filenames or relative paths.
+
+If you trained a model with the main
+[`dusty-lm`](https://github.com/khordoo/dusty-lm) repository, its checkpoint and
+tokenizer are stored in separate artifact directories. Load the promoted SFT
+checkpoint with:
+
+```python
+from dustylm import DustyLM
+
+model = DustyLM.from_pretrained(
+    "artifacts/checkpoints",
+    model_file="dusty8m_sft.pt",
+    tokenizer_file="../tokenizers/dusty_tokenizer.json",
+)
+
+response = model.chat([{"role": "user", "content": "who are you?"}])
+```
+
+To inspect an intermediate checkpoint instead, change `model_file` to an
+existing step file such as `dusty8m_sft_step_100.pt`.
+
+The SDK infers several architecture dimensions from PyTorch checkpoint shapes
+and currently expects DustyLM-compatible configurations.
+
+## API
+
+### DustyLM.from_pretrained
+
+```python
+DustyLM.from_pretrained(
+    repo_id_or_path: str = "mkhordoo/dusty-8m-sft",
+    *,
+    model_file: str | None = None,
+    tokenizer_file: str | None = None,
+    backend: str = "torch",
+)
+```
+
+| Argument | Default | Description |
+|---|---|---|
+| `repo_id_or_path` | `"mkhordoo/dusty-8m-sft"` | HF Hub repo ID or local directory path |
+| `model_file` | `"model.pt"` (torch) / `"model_int8.onnx"` (onnx) | Override the model filename in the directory |
+| `tokenizer_file` | `"tokenizer.json"` | Override the tokenizer filename |
+| `backend` | `"torch"` | `"torch"` or `"onnx"` |
+
+### DustyLM.chat
+
+```python
+model.chat(
+    messages: list[dict],
+    temperature: float = 0.7,
+    max_tokens: int = 64,
+    top_p: float = 0.9,
+) -> dict
+```
+
+Returns an OpenAI-style chat completion dict.
+
+---
+
+Built from the [dusty-lm](https://github.com/khordoo/dusty-lm) training repository.

dustylm_sdk-0.1.1/README.md

@@ -0,0 +1,118 @@
+<p align="center">
+  <img src="https://raw.githubusercontent.com/khordoo/dustylm-sdk/main/assets/logo.png" alt="DustyLM Logo" width="200"/>
+</p>
+
+# dustylm
+
+Run DustyLM, the 8M parameter robot vacuum language model, in a few lines.
+
+```
+pip install dustylm-sdk
+```
+
+```python
+from dustylm import DustyLM
+model = DustyLM.from_pretrained("mkhordoo/dusty-8m-sft")
+response = model.chat([{"role": "user", "content": "who are you?"}])
+print(response["choices"][0]["message"]["content"])
+
+# beep. i am a little robot. i clean floors and find crumbs.
+```
+
+## Backends
+
+| Backend | Install | Model file | Best for |
+|---|---|---:|---|
+| `torch` (default) | `pip install dustylm-sdk` | ~32 MB (FP32) | Inspection and experimentation |
+| `onnx` | `pip install dustylm-sdk[onnx]` | ~8–10 MB (int8) | A compact model artifact |
+
+These sizes describe the model files, not the complete Python environment.
+Installation size varies by platform. Both backends currently include PyTorch as
+an SDK dependency.
+
+```python
+# ONNX backend
+model = DustyLM.from_pretrained("mkhordoo/dusty-8m-sft", backend="onnx")
+```
+
+## Loading Local Checkpoints
+
+Pass a Hugging Face repository ID to download its model and tokenizer
+automatically, as shown in the quick-start example:
+
+```python
+model = DustyLM.from_pretrained("mkhordoo/dusty-8m-sft")
+```
+
+For files already downloaded to your machine, pass their local directory
+instead. With the default filenames, that directory must contain
+`tokenizer.json` and either `model.pt` (PyTorch) or `model_int8.onnx` (ONNX):
+
+```python
+model = DustyLM.from_pretrained("./my-checkpoint/")
+```
+
+You can use any local directory as the artifact root. `model_file` and
+`tokenizer_file` are resolved relative to that directory, so they may use
+different filenames or relative paths.
+
+If you trained a model with the main
+[`dusty-lm`](https://github.com/khordoo/dusty-lm) repository, its checkpoint and
+tokenizer are stored in separate artifact directories. Load the promoted SFT
+checkpoint with:
+
+```python
+from dustylm import DustyLM
+
+model = DustyLM.from_pretrained(
+    "artifacts/checkpoints",
+    model_file="dusty8m_sft.pt",
+    tokenizer_file="../tokenizers/dusty_tokenizer.json",
+)
+
+response = model.chat([{"role": "user", "content": "who are you?"}])
+```
+
+To inspect an intermediate checkpoint instead, change `model_file` to an
+existing step file such as `dusty8m_sft_step_100.pt`.
+
+The SDK infers several architecture dimensions from PyTorch checkpoint shapes
+and currently expects DustyLM-compatible configurations.
+
+## API
+
+### DustyLM.from_pretrained
+
+```python
+DustyLM.from_pretrained(
+    repo_id_or_path: str = "mkhordoo/dusty-8m-sft",
+    *,
+    model_file: str | None = None,
+    tokenizer_file: str | None = None,
+    backend: str = "torch",
+)
+```
+
+| Argument | Default | Description |
+|---|---|---|
+| `repo_id_or_path` | `"mkhordoo/dusty-8m-sft"` | HF Hub repo ID or local directory path |
+| `model_file` | `"model.pt"` (torch) / `"model_int8.onnx"` (onnx) | Override the model filename in the directory |
+| `tokenizer_file` | `"tokenizer.json"` | Override the tokenizer filename |
+| `backend` | `"torch"` | `"torch"` or `"onnx"` |
+
+### DustyLM.chat
+
+```python
+model.chat(
+    messages: list[dict],
+    temperature: float = 0.7,
+    max_tokens: int = 64,
+    top_p: float = 0.9,
+) -> dict
+```
+
+Returns an OpenAI-style chat completion dict.
+
+---
+
+Built from the [dusty-lm](https://github.com/khordoo/dusty-lm) training repository.

dustylm_sdk-0.1.1/dustylm/__init__.py

@@ -0,0 +1,4 @@
+from dustylm.inference import DustyLM
+
+__all__ = ["DustyLM"]
+__version__ = "0.1.1"

dustylm_sdk-0.1.1/dustylm/inference.py

@@ -0,0 +1,461 @@
+"""Inference-only DustyLM — load from Hugging Face Hub or a local directory."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal
+
+import torch
+import torch.nn as nn
+
+Backend = Literal["torch", "onnx"]
+
+HUB_REPO = "mkhordoo/dusty-8m-sft"
+HUB_FILES = {
+    "torch": "model.pt",
+    "onnx": "model_int8.onnx",
+    "tokenizer": "tokenizer.json",
+}
+
+CHATML_START_TOKEN = "<|im_start|>"
+CHATML_END_TOKEN = "<|im_end|>"
+SUPPORTED_ROLES = {"system", "user", "assistant"}
+
+DEFAULT_CONFIG = {
+    "num_layers": 8,
+    "vocab_size": 4096,
+    "max_seq_len": 256,
+    "embed_dim": 256,
+    "num_heads": 8,
+    "num_kv_heads": 4,
+    "hidden_dim": 1024,
+}
+
+
+# ── Model (PyTorch) ──────────────────────────────────────────────────────────
+
+
+def model_from_config(config: dict) -> "DustyLMModel":
+    model = DustyLMModel(**config)
+    return model
+
+
+@dataclass(frozen=True)
+class ForwardContext:
+    position_ids: "torch.Tensor"
+    rope_sin: "torch.Tensor"
+    rope_cos: "torch.Tensor"
+
+
+def _rotate_half(x):
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat([-x2, x1], dim=-1)
+
+
+def _apply_rope(x, sin, cos):
+    return x * cos + _rotate_half(x) * sin
+
+
+class RotaryPositionalEmbedding(nn.Module):
+    def __init__(self, head_dim: int, max_seq_len: int, base: int = 10000):
+        super().__init__()
+        inv_freq = base ** (-torch.arange(0, head_dim, 2) / head_dim)
+        inv_freq = torch.cat([inv_freq, inv_freq], dim=-1)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self._build_cache(max_seq_len)
+
+    def _build_cache(self, max_seq_len: int):
+        positions = torch.arange(
+            max_seq_len, dtype=self.inv_freq.dtype, device=self.inv_freq.device
+        )
+        emb = torch.outer(positions, self.inv_freq)
+        self.register_buffer("sin_cache", emb.sin(), persistent=False)
+        self.register_buffer("cos_cache", emb.cos(), persistent=False)
+
+    def forward(self, position_ids):
+        sin = self.sin_cache[position_ids].unsqueeze(0).unsqueeze(1)
+        cos = self.cos_cache[position_ids].unsqueeze(0).unsqueeze(1)
+        return sin, cos
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads, num_kv_heads, max_seq_len):
+        super().__init__()
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = embed_dim // num_heads
+        self.kv_repeat = num_heads // num_kv_heads
+        self.qkv_proj = nn.Linear(
+            embed_dim, embed_dim + 2 * num_kv_heads * self.head_dim, bias=False
+        )
+        self.proj = nn.Linear(embed_dim, embed_dim, bias=False)
+
+    def forward(self, x, context, past_kv=None):
+        B, T, D = x.shape
+        qkv = self.qkv_proj(x)
+        q_dim = self.num_heads * self.head_dim
+        kv_dim = self.num_kv_heads * self.head_dim
+        q, k, v = qkv.split([q_dim, kv_dim, kv_dim], dim=-1)
+
+        q = q.reshape(B, T, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.reshape(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = v.reshape(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2)
+
+        q = _apply_rope(q, context.rope_sin, context.rope_cos)
+        k = _apply_rope(k, context.rope_sin, context.rope_cos)
+
+        past_len = 0
+        if past_kv is not None:
+            pk, pv = past_kv
+            past_len = pk.shape[2]
+            k = torch.cat([pk, k], dim=2)
+            v = torch.cat([pv, v], dim=2)
+
+        present_kv = (k, v)
+        k = k.repeat_interleave(self.kv_repeat, dim=1)
+        v = v.repeat_interleave(self.kv_repeat, dim=1)
+
+        scores = q @ k.transpose(-2, -1) / (self.head_dim**0.5)
+        if T > 1:
+            total_len = past_len + T
+            mask = torch.triu(
+                torch.ones(T, total_len, device=x.device, dtype=torch.bool),
+                diagonal=past_len + 1,
+            )
+            scores = scores.masked_fill(mask, float("-inf"))
+
+        attn = torch.softmax(scores, dim=-1)
+        out = (attn @ v).transpose(1, 2).reshape(B, T, D)
+        return self.proj(out), present_kv
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self, embed_dim, num_heads, num_kv_heads, max_seq_len, hidden_dim=None
+    ):
+        super().__init__()
+        hidden_dim = hidden_dim or 4 * embed_dim
+        self.att_norm = nn.RMSNorm(embed_dim)
+        self.attention = MultiHeadAttention(
+            embed_dim, num_heads, num_kv_heads, max_seq_len
+        )
+        self.mlp = nn.Sequential(
+            nn.Linear(embed_dim, hidden_dim, bias=False),
+            nn.GELU(),
+            nn.Linear(hidden_dim, embed_dim, bias=False),
+        )
+        self.mlp_norm = nn.RMSNorm(embed_dim)
+
+    def forward(self, x, context, kv_cache=None):
+        a, pkv = self.attention(self.att_norm(x), context, kv_cache)
+        x = x + a
+        x = x + self.mlp(self.mlp_norm(x))
+        return x, pkv
+
+
+class DustyLMModel(nn.Module):
+    def __init__(
+        self,
+        num_layers,
+        vocab_size,
+        max_seq_len,
+        embed_dim,
+        num_heads,
+        num_kv_heads,
+        hidden_dim=None,
+        rope_base=10000,
+    ):
+        super().__init__()
+        head_dim = embed_dim // num_heads
+        self.rope = RotaryPositionalEmbedding(head_dim, max_seq_len, rope_base)
+        self.embed = nn.Embedding(vocab_size, embed_dim)
+        self.layers = nn.ModuleList(
+            [
+                TransformerBlock(
+                    embed_dim, num_heads, num_kv_heads, max_seq_len, hidden_dim
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.final_norm = nn.RMSNorm(embed_dim)
+        self.vocab_proj = nn.Linear(embed_dim, vocab_size, bias=False)
+        self.max_seq_len = max_seq_len
+
+    def forward(self, x, kv_cache=None):
+        B, T = x.shape
+        use_cache = kv_cache is not None
+
+        past_len = 0
+        if use_cache and kv_cache[0] is not None:
+            past_len = kv_cache[0][0].shape[2]
+
+        position_ids = torch.arange(
+            past_len, past_len + T, dtype=torch.long, device=x.device
+        )
+        sin, cos = self.rope(position_ids)
+        ctx = ForwardContext(position_ids=position_ids, rope_sin=sin, rope_cos=cos)
+
+        h = self.embed(x)
+        next_kv = [] if use_cache else None
+        for i, layer in enumerate(self.layers):
+            h, pkv = layer(h, ctx, kv_cache[i] if use_cache else None)
+            if use_cache:
+                next_kv.append(pkv)
+
+        logits = self.vocab_proj(self.final_norm(h))
+        if use_cache:
+            return logits, next_kv
+        return logits
+
+
+# ── Tokenizer ─────────────────────────────────────────────────────────────────
+
+
+def load_tokenizer(path: Path):
+    from tokenizers import Tokenizer
+
+    return Tokenizer.from_file(str(path))
+
+
+def format_chatml(messages: list[dict]) -> str:
+    chunks = []
+    for msg in messages:
+        role, content = msg.get("role", "user"), msg.get("content", "")
+        if not content.strip():
+            continue
+        chunks.append(f"{CHATML_START_TOKEN}{role}\n{content}{CHATML_END_TOKEN}\n")
+    chunks.append(f"{CHATML_START_TOKEN}assistant\n")
+    return "".join(chunks)
+
+
+# ── Sampling ──────────────────────────────────────────────────────────────────
+
+
+def sample(logits, temperature=0.7, top_p=0.9):
+    logits = logits / max(temperature, 1e-8)
+    probs = torch.softmax(logits, dim=-1)
+
+    if top_p < 1.0:
+        sorted_probs, sorted_indices = torch.sort(probs, descending=True)
+        cumsum = torch.cumsum(sorted_probs, dim=-1)
+        remove = cumsum > top_p
+        remove[..., 1:] = remove[..., :-1].clone()
+        remove[..., 0] = False
+        sorted_probs[remove] = 0.0
+        probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)
+        probs = probs.gather(-1, sorted_indices.argsort(-1))
+
+    return int(torch.multinomial(probs, 1).item())
+
+
+# ── Main class ───────────────────────────────────────────────────────────────
+
+
+class DustyLM:
+    """Load and chat with a DustyLM model."""
+
+    def __init__(self, model, tokenizer, config: dict, backend: Backend):
+        self._model = model
+        self._tokenizer = tokenizer
+        self._config = config
+        self._backend = backend
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        repo_id_or_path: str = HUB_REPO,
+        *,
+        model_file: str | None = None,
+        tokenizer_file: str | None = None,
+        backend: Backend = "torch",
+    ) -> "DustyLM":
+        """Load from Hugging Face Hub or a local directory.
+
+        Args:
+            repo_id_or_path: HF repo ID (e.g. "mkhordoo/dusty-8m-sft")
+                             or local directory path.
+            model_file: Model filename in the directory (default depends on backend:
+                        "model.pt" for torch, "model_int8.onnx" for onnx).
+            tokenizer_file: Tokenizer filename (default "tokenizer.json").
+            backend: "torch" (default) or "onnx".
+        """
+        model_file = model_file or HUB_FILES[backend if backend == "onnx" else "torch"]
+        tokenizer_file = tokenizer_file or HUB_FILES["tokenizer"]
+
+        path = Path(repo_id_or_path)
+        is_local = path.exists()
+
+        if is_local:
+            tokenizer_path = path / tokenizer_file
+        else:
+            from huggingface_hub import hf_hub_download
+
+            tokenizer_path = Path(
+                hf_hub_download(repo_id=repo_id_or_path, filename=tokenizer_file)
+            )
+
+        tokenizer = load_tokenizer(tokenizer_path)
+
+        if backend == "torch":
+            model, config = cls._load_torch(repo_id_or_path, model_file, is_local)
+        elif backend == "onnx":
+            model, config = cls._load_onnx(repo_id_or_path, model_file, is_local)
+        else:
+            raise ValueError(f"Unknown backend: {backend}")
+
+        return cls(model=model, tokenizer=tokenizer, config=config, backend=backend)
+
+    @classmethod
+    def _load_torch(cls, repo_id_or_path, model_file, is_local):
+        if is_local:
+            ckpt_path = Path(repo_id_or_path) / model_file
+        else:
+            from huggingface_hub import hf_hub_download
+
+            ckpt_path = Path(
+                hf_hub_download(repo_id=repo_id_or_path, filename=model_file)
+            )
+
+        state = torch.load(ckpt_path, map_location="cpu", weights_only=True)
+        state.pop("rope.sin_cache", None)
+        state.pop("rope.cos_cache", None)
+
+        config = dict(DEFAULT_CONFIG)
+        config["vocab_size"] = state["vocab_proj.weight"].shape[0]
+        config["embed_dim"] = state["embed.weight"].shape[1]
+        config["num_layers"] = sum(
+            1
+            for k in state
+            if k.startswith("layers.") and k.endswith(".att_norm.weight")
+        )
+        qkv_out = state["layers.0.attention.qkv_proj.weight"].shape[0]
+        head_dim = config["embed_dim"] // config["num_heads"]
+        num_kv_heads = (qkv_out - config["embed_dim"]) // (2 * head_dim)
+        config["num_kv_heads"] = num_kv_heads
+        config["num_heads"] = config["num_heads"] or config["num_kv_heads"]
+        config["hidden_dim"] = state["layers.0.mlp.0.weight"].shape[0]
+
+        model = model_from_config(config)
+        model.load_state_dict(state)
+        model.eval()
+        return model, config
+
+    @classmethod
+    def _load_onnx(cls, repo_id_or_path, model_file, is_local):
+        import onnxruntime as ort
+
+        if is_local:
+            onnx_path = str(Path(repo_id_or_path) / model_file)
+        else:
+            from huggingface_hub import hf_hub_download
+
+            onnx_path = hf_hub_download(repo_id=repo_id_or_path, filename=model_file)
+
+        session = ort.InferenceSession(onnx_path)
+        config = dict(DEFAULT_CONFIG)
+        return session, config
+
+    def chat(
+        self,
+        messages: list[dict[str, Any]],
+        temperature: float = 0.7,
+        max_tokens: int = 64,
+        top_p: float = 0.9,
+    ) -> dict:
+        """Chat completion returning an OpenAI-style response dict."""
+        prompt = format_chatml(messages)
+
+        if self._backend == "torch":
+            return self._chat_torch(prompt, temperature, max_tokens, top_p)
+        else:
+            return self._chat_onnx(prompt, temperature, max_tokens, top_p)
+
+    def _chat_torch(self, prompt, temperature, max_tokens, top_p):
+        prompt_ids = self._tokenizer.encode(prompt).ids
+        max_seq = self._config["max_seq_len"]
+
+        input_ids = torch.tensor([prompt_ids])
+        past_len = input_ids.shape[1]
+
+        if max_seq - past_len < max_tokens:
+            max_tokens = max_seq - past_len
+
+        if max_tokens <= 0:
+            return self._response("", 0)
+
+        model = self._model
+        device = next(model.parameters()).device
+        input_ids = input_ids.to(device)
+
+        num_layers = self._config["num_layers"]
+        kv_cache = [None] * num_layers
+        logits, kv_cache = model(input_ids, kv_cache)
+        end_token_id = self._tokenizer.token_to_id(CHATML_END_TOKEN)
+        generated = []
+
+        for _ in range(max_tokens):
+            next_id = sample(logits[0, -1, :], temperature, top_p)
+            generated.append(next_id)
+            if next_id == end_token_id:
+                break
+
+            inp = torch.tensor([[next_id]], device=device)
+            logits, kv_cache = model(inp, kv_cache)
+
+        content = self._decode_completion(generated)
+        return self._response(content, len(generated))
+
+    def _chat_onnx(self, prompt, temperature, max_tokens, top_p):
+        import numpy as np
+
+        prompt_ids = self._tokenizer.encode(prompt).ids
+        max_seq = self._config["max_seq_len"]
+
+        gen_count = min(max_tokens, max_seq - len(prompt_ids))
+        if gen_count <= 0:
+            content = self._tokenizer.decode(prompt_ids).strip()
+            return self._response(content, 0)
+
+        session = self._model
+        end_token_id = self._tokenizer.token_to_id(CHATML_END_TOKEN)
+        generated = []
+
+        for step in range(gen_count):
+            input_ids = np.array([prompt_ids + generated], dtype=np.int64)
+            outputs = session.run(None, {"input_ids": input_ids})
+            logits = outputs[0]
+            last_logits = torch.tensor(logits[0, -1, :])
+            next_id = sample(last_logits, temperature, top_p)
+            generated.append(next_id)
+            if next_id == end_token_id:
+                break
+
+        content = self._decode_completion(generated)
+        return self._response(content, len(generated))
+
+    def _decode_completion(self, token_ids: list[int]) -> str:
+        end_token_id = self._tokenizer.token_to_id(CHATML_END_TOKEN)
+        if end_token_id in token_ids:
+            token_ids = token_ids[: token_ids.index(end_token_id)]
+        return self._tokenizer.decode(token_ids).strip()
+
+    @staticmethod
+    def _response(content: str, completion_tokens: int) -> dict:
+        return {
+            "object": "chat.completion",
+            "model": "dusty8m",
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": content},
+                    "finish_reason": "stop" if completion_tokens > 0 else "length",
+                }
+            ],
+            "usage": {
+                "prompt_tokens": 0,
+                "completion_tokens": completion_tokens,
+                "total_tokens": completion_tokens,
+            },
+        }

dustylm_sdk-0.1.1/dustylm_sdk.egg-info/PKG-INFO

@@ -0,0 +1,143 @@
+Metadata-Version: 2.4
+Name: dustylm-sdk
+Version: 0.1.1
+Summary: The official Python SDK for DustyLM: an 8M-parameter model that talks like a robot vacuum.
+Author-email: Mahmood Khordoo <m.khordoo@gmail.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/khordoo/dustylm-sdk
+Project-URL: Source, https://github.com/khordoo/dustylm-sdk
+Keywords: llm,slm,small-language-model,transformer,pytorch,onnx,inference
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0
+Requires-Dist: tokenizers>=0.15
+Requires-Dist: huggingface-hub>=0.20
+Requires-Dist: numpy>=1.24
+Provides-Extra: onnx
+Requires-Dist: onnxruntime>=1.15; extra == "onnx"
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/khordoo/dustylm-sdk/main/assets/logo.png" alt="DustyLM Logo" width="200"/>
+</p>
+
+# dustylm
+
+Run DustyLM, the 8M parameter robot vacuum language model, in a few lines.
+
+```
+pip install dustylm-sdk
+```
+
+```python
+from dustylm import DustyLM
+model = DustyLM.from_pretrained("mkhordoo/dusty-8m-sft")
+response = model.chat([{"role": "user", "content": "who are you?"}])
+print(response["choices"][0]["message"]["content"])
+
+# beep. i am a little robot. i clean floors and find crumbs.
+```
+
+## Backends
+
+| Backend | Install | Model file | Best for |
+|---|---|---:|---|
+| `torch` (default) | `pip install dustylm-sdk` | ~32 MB (FP32) | Inspection and experimentation |
+| `onnx` | `pip install dustylm-sdk[onnx]` | ~8–10 MB (int8) | A compact model artifact |
+
+These sizes describe the model files, not the complete Python environment.
+Installation size varies by platform. Both backends currently include PyTorch as
+an SDK dependency.
+
+```python
+# ONNX backend
+model = DustyLM.from_pretrained("mkhordoo/dusty-8m-sft", backend="onnx")
+```
+
+## Loading Local Checkpoints
+
+Pass a Hugging Face repository ID to download its model and tokenizer
+automatically, as shown in the quick-start example:
+
+```python
+model = DustyLM.from_pretrained("mkhordoo/dusty-8m-sft")
+```
+
+For files already downloaded to your machine, pass their local directory
+instead. With the default filenames, that directory must contain
+`tokenizer.json` and either `model.pt` (PyTorch) or `model_int8.onnx` (ONNX):
+
+```python
+model = DustyLM.from_pretrained("./my-checkpoint/")
+```
+
+You can use any local directory as the artifact root. `model_file` and
+`tokenizer_file` are resolved relative to that directory, so they may use
+different filenames or relative paths.
+
+If you trained a model with the main
+[`dusty-lm`](https://github.com/khordoo/dusty-lm) repository, its checkpoint and
+tokenizer are stored in separate artifact directories. Load the promoted SFT
+checkpoint with:
+
+```python
+from dustylm import DustyLM
+
+model = DustyLM.from_pretrained(
+    "artifacts/checkpoints",
+    model_file="dusty8m_sft.pt",
+    tokenizer_file="../tokenizers/dusty_tokenizer.json",
+)
+
+response = model.chat([{"role": "user", "content": "who are you?"}])
+```
+
+To inspect an intermediate checkpoint instead, change `model_file` to an
+existing step file such as `dusty8m_sft_step_100.pt`.
+
+The SDK infers several architecture dimensions from PyTorch checkpoint shapes
+and currently expects DustyLM-compatible configurations.
+
+## API
+
+### DustyLM.from_pretrained
+
+```python
+DustyLM.from_pretrained(
+    repo_id_or_path: str = "mkhordoo/dusty-8m-sft",
+    *,
+    model_file: str | None = None,
+    tokenizer_file: str | None = None,
+    backend: str = "torch",
+)
+```
+
+| Argument | Default | Description |
+|---|---|---|
+| `repo_id_or_path` | `"mkhordoo/dusty-8m-sft"` | HF Hub repo ID or local directory path |
+| `model_file` | `"model.pt"` (torch) / `"model_int8.onnx"` (onnx) | Override the model filename in the directory |
+| `tokenizer_file` | `"tokenizer.json"` | Override the tokenizer filename |
+| `backend` | `"torch"` | `"torch"` or `"onnx"` |
+
+### DustyLM.chat
+
+```python
+model.chat(
+    messages: list[dict],
+    temperature: float = 0.7,
+    max_tokens: int = 64,
+    top_p: float = 0.9,
+) -> dict
+```
+
+Returns an OpenAI-style chat completion dict.
+
+---
+
+Built from the [dusty-lm](https://github.com/khordoo/dusty-lm) training repository.

dustylm_sdk-0.1.1/dustylm_sdk.egg-info/SOURCES.txt

@@ -0,0 +1,9 @@
+README.md
+pyproject.toml
+dustylm/__init__.py
+dustylm/inference.py
+dustylm_sdk.egg-info/PKG-INFO
+dustylm_sdk.egg-info/SOURCES.txt
+dustylm_sdk.egg-info/dependency_links.txt
+dustylm_sdk.egg-info/requires.txt
+dustylm_sdk.egg-info/top_level.txt

dustylm_sdk-0.1.1/dustylm_sdk.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

dustylm_sdk-0.1.1/dustylm_sdk.egg-info/requires.txt

@@ -0,0 +1,7 @@
+torch>=2.0
+tokenizers>=0.15
+huggingface-hub>=0.20
+numpy>=1.24
+
+[onnx]
+onnxruntime>=1.15

dustylm_sdk-0.1.1/dustylm_sdk.egg-info/top_level.txt

@@ -0,0 +1 @@
+dustylm

dustylm_sdk-0.1.1/pyproject.toml

@@ -0,0 +1,50 @@
+[build-system]
+requires = ["setuptools>=64", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "dustylm-sdk"
+version = "0.1.1"
+description = "The official Python SDK for DustyLM: an 8M-parameter model that talks like a robot vacuum."
+keywords = [
+    "llm",
+    "slm",
+    "small-language-model",
+    "transformer",
+    "pytorch",
+    "onnx",
+    "inference"
+]
+readme = "README.md"
+license = "MIT"
+authors = [
+    { name = "Mahmood Khordoo", email = "m.khordoo@gmail.com" },
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Education",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+requires-python = ">=3.11"
+dependencies = [
+    "torch>=2.0",
+    "tokenizers>=0.15",
+    "huggingface-hub>=0.20",
+    "numpy>=1.24",
+]
+
+[project.optional-dependencies]
+onnx = [
+    "onnxruntime>=1.15",
+]
+
+[project.urls]
+Homepage = "https://github.com/khordoo/dustylm-sdk"
+Source = "https://github.com/khordoo/dustylm-sdk"
+
+[tool.setuptools.packages.find]
+include = ["dustylm*"]

dustylm_sdk-0.1.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+