@tuned-tensor/local 0.2.9 → 0.4.0

package/training/local-runner/src/serve.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-import json
-import base64
 import hmac
+import json
 import math
 import os
 import sys
@@ -14,21 +13,16 @@ from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import Any
 
-# HF_HOME is supplied by the Node launcher before this module starts. Importing
-# the model stack only after reading process configuration keeps every local
-# workflow on the same Hugging Face cache contract.
-from evaluate import (  # noqa: E402
+from evaluate import (
     configure_hugging_face_cache,
-    generate_multimodal_one,
     import_runtime_dependencies,
-    load_multimodal_model,
     load_text_model,
     resolve_adapter_path,
     sampling_kwargs,
 )
 
 
-MODEL_ARTIFACT = os.environ["TT_MODEL_ARTIFACT"]
+MODEL_ARTIFACT = os.environ.get("TT_MODEL_ARTIFACT")
 BASE_MODEL = os.environ["TT_BASE_MODEL"]
 BASE_MODEL_REVISION = os.environ.get("TT_BASE_MODEL_REVISION")
 MODEL_NAME = os.environ.get("TT_MODEL_NAME", "tuned-tensor-local")
@@ -36,117 +30,99 @@ MODEL_LOADER = os.environ.get("TT_MODEL_LOADER", "causal_lm")
 SYSTEM_PROMPT = os.environ.get("TT_SYSTEM_PROMPT", "").strip()
 HOST = os.environ.get("TT_HOST", "127.0.0.1")
 PORT = int(os.environ.get("TT_PORT", "8000"))
-DEVICE_REQUEST = os.environ.get("TT_DEVICE", "auto")
-TRUST_REMOTE_CODE = os.environ.get("TT_TRUST_REMOTE_CODE", "true").lower() in {"1", "true", "yes"}
+DEVICE_REQUEST = os.environ.get("TT_DEVICE", "cuda")
 DEFAULT_MAX_TOKENS = int(os.environ.get("TT_MAX_TOKENS", "512"))
 DEFAULT_TEMPERATURE = float(os.environ.get("TT_TEMPERATURE", "0"))
 DEFAULT_TOP_P = float(os.environ.get("TT_TOP_P", "1"))
-CHAT_TEMPLATE_KWARGS = json.loads(os.environ.get("TT_CHAT_TEMPLATE_KWARGS", "{}"))
 MAX_REQUEST_BYTES = 2 * 1024 * 1024
 MAX_PROMPT_CHARS = 100_000
 MAX_PROMPT_TOKENS = 16_384
 MAX_MESSAGES = 128
-MAX_IMAGE_BYTES = 5 * 1024 * 1024
 API_KEY = os.environ.get("TT_API_KEY", "")
 MAX_CONCURRENT_REQUESTS = int(os.environ.get("TT_MAX_CONCURRENT_REQUESTS", "1"))
 
 
+if MODEL_LOADER != "causal_lm":
+    raise ValueError("The bundled model server is text-only and requires TT_MODEL_LOADER=causal_lm")
 configure_hugging_face_cache(os.environ.get("HF_HOME"))
 import_runtime_dependencies()
 TEMP_DIR = TemporaryDirectory(prefix="tt-local-serve-")
-ADAPTER_PATH = resolve_adapter_path(MODEL_ARTIFACT, Path(TEMP_DIR.name))
+ADAPTER_PATH = (
+    resolve_adapter_path(MODEL_ARTIFACT, Path(TEMP_DIR.name))
+    if MODEL_ARTIFACT
+    else None
+)
 MODEL_PAYLOAD = {
     "base_model": BASE_MODEL,
     "base_model_revision": BASE_MODEL_REVISION,
     "device": DEVICE_REQUEST,
-    "trust_remote_code": TRUST_REMOTE_CODE,
     "model_loader": MODEL_LOADER,
 }
-
-if MODEL_LOADER == "image_text_to_text":
-    MODEL, PROCESSOR, DEVICE = load_multimodal_model(MODEL_PAYLOAD, ADAPTER_PATH)
-    TOKENIZER = getattr(PROCESSOR, "tokenizer", PROCESSOR)
-else:
-    MODEL, TOKENIZER, DEVICE = load_text_model(MODEL_PAYLOAD, ADAPTER_PATH)
-    PROCESSOR = None
-
+MODEL, TOKENIZER, DEVICE = load_text_model(MODEL_PAYLOAD, ADAPTER_PATH)
 GENERATION_LOCK = threading.Lock()
 REQUEST_SLOTS = threading.BoundedSemaphore(MAX_CONCURRENT_REQUESTS)
 
 
-def normalize_messages(raw_messages: Any) -> list[dict[str, Any]]:
-    if not isinstance(raw_messages, list) or not raw_messages:
-        raise ValueError("Request must include a non-empty messages array.")
-    if len(raw_messages) > MAX_MESSAGES:
-        raise ValueError(f"Request exceeds the {MAX_MESSAGES}-message limit.")
-    messages: list[dict[str, Any]] = []
-    for raw in raw_messages:
-        if not isinstance(raw, dict) or raw.get("role") not in {"system", "user", "assistant", "tool"}:
-            raise ValueError("Each message must have a supported role and content.")
-        messages.append({"role": raw["role"], "content": raw.get("content", "")})
-    # A spec passed by the server owner is invariant behavior. Client-supplied
-    # system messages may add context but cannot replace it.
-    if SYSTEM_PROMPT:
-        messages.insert(0, {"role": "system", "content": SYSTEM_PROMPT})
-    if sum(len(text_content(message.get("content"))) for message in messages) > MAX_PROMPT_CHARS:
-        raise ValueError(f"Prompt exceeds the {MAX_PROMPT_CHARS}-character limit.")
-    return messages
-
-
 def text_content(content: Any) -> str:
     if isinstance(content, str):
         return content
+    if content is None:
+        return ""
     if not isinstance(content, list):
-        return str(content) if content is not None else ""
+        raise ValueError("Message content must be text or an array of text parts.")
     parts: list[str] = []
     for part in content:
-        if isinstance(part, dict) and part.get("type") == "text":
-            parts.append(str(part.get("text", "")))
-        elif isinstance(part, dict) and part.get("type") in {"image", "image_url"}:
-            continue
-        else:
-            parts.append(str(part))
-    return "\n".join(part for part in parts if part)
+        if not isinstance(part, dict) or part.get("type") != "text":
+            raise ValueError("The bundled Qwen model server accepts text content only.")
+        text = part.get("text")
+        if not isinstance(text, str):
+            raise ValueError("Every text content part must contain a string text field.")
+        parts.append(text)
+    return "\n".join(parts)
 
 
-def image_assets(messages: list[dict[str, Any]]) -> list[dict[str, str]]:
-    assets: list[dict[str, str]] = []
-    for message in messages:
-        content = message.get("content")
-        if not isinstance(content, list):
-            continue
-        for part in content:
-            if not isinstance(part, dict) or part.get("type") not in {"image", "image_url"}:
-                continue
-            image = part.get("image_url")
-            if isinstance(image, dict):
-                image = image.get("url")
-            image = image or part.get("image") or part.get("data_uri") or part.get("uri") or part.get("path")
-            if isinstance(image, str) and image:
-                if not image.startswith("data:image/") or "," not in image:
-                    raise ValueError("Serving accepts image content only as data:image/... URIs.")
-                encoded = image.split(",", 1)[1]
-                try:
-                    decoded_size = len(base64.b64decode(encoded, validate=True))
-                except Exception as exc:
-                    raise ValueError("Image data URI contains invalid base64.") from exc
-                if decoded_size > MAX_IMAGE_BYTES:
-                    raise ValueError(f"Image exceeds the {MAX_IMAGE_BYTES}-byte decoded limit.")
-                assets.append({"type": "image", "image": image})
-    return assets
+def normalize_messages(raw_messages: Any) -> list[dict[str, str]]:
+    if not isinstance(raw_messages, list) or not raw_messages:
+        raise ValueError("Request must include a non-empty messages array.")
+    if len(raw_messages) > MAX_MESSAGES:
+        raise ValueError(f"Request exceeds the {MAX_MESSAGES}-message limit.")
 
+    system_parts = [SYSTEM_PROMPT] if SYSTEM_PROMPT else []
+    conversation: list[dict[str, str]] = []
+    for raw in raw_messages:
+        if not isinstance(raw, dict) or raw.get("role") not in {"system", "user", "assistant", "tool"}:
+            raise ValueError("Each message must have a supported role and text content.")
+        role = str(raw["role"])
+        content = text_content(raw.get("content", ""))
+        if role == "system":
+            if content.strip():
+                system_parts.append(content.strip())
+        else:
+            conversation.append({"role": role, "content": content})
+
+    if not conversation:
+        raise ValueError("Request must contain at least one non-system message.")
+    messages: list[dict[str, str]] = []
+    if system_parts:
+        # Qwen permits one leading system message. Merge the invariant owner
+        # prompt and any client context instead of triggering an implicit
+        # template fallback with duplicate system turns.
+        messages.append({"role": "system", "content": "\n\n".join(system_parts)})
+    messages.extend(conversation)
+    if sum(len(message["content"]) for message in messages) > MAX_PROMPT_CHARS:
+        raise ValueError(f"Prompt exceeds the {MAX_PROMPT_CHARS}-character limit.")
+    return messages
 
-def generate_text(messages: list[dict[str, Any]], generation: dict[str, Any]) -> tuple[str, int, int]:
-    normalized = [{"role": message["role"], "content": text_content(message.get("content"))} for message in messages]
+
+def generate_text(messages: list[dict[str, str]], generation: dict[str, Any]) -> tuple[str, int, int]:
     try:
         prompt = TOKENIZER.apply_chat_template(
-            normalized,
+            messages,
             tokenize=False,
             add_generation_prompt=True,
-            **CHAT_TEMPLATE_KWARGS,
         )
-    except Exception:
-        prompt = "\n".join(f"{message['role']}: {message['content']}" for message in normalized) + "\nassistant:"
+    except Exception as exc:
+        raise ValueError(f"Qwen chat-template rendering failed: {exc}") from exc
     inputs = TOKENIZER(prompt, return_tensors="pt")
     if int(inputs["input_ids"].shape[-1]) > MAX_PROMPT_TOKENS:
         raise ValueError(f"Prompt exceeds the {MAX_PROMPT_TOKENS}-token limit.")
@@ -155,7 +131,7 @@ def generate_text(messages: list[dict[str, Any]], generation: dict[str, Any]) ->
     with GENERATION_LOCK:
         import torch
 
-        with torch.no_grad():
+        with torch.inference_mode():
             generated = MODEL.generate(
                 **inputs,
                 max_new_tokens=int(generation["max_new_tokens"]),
@@ -168,33 +144,6 @@ def generate_text(messages: list[dict[str, Any]], generation: dict[str, Any]) ->
     return content, prompt_tokens, int(completion_tokens.shape[-1])
 
 
-def generate_multimodal(messages: list[dict[str, Any]], generation: dict[str, Any]) -> tuple[str, int, int]:
-    system_parts = [text_content(message.get("content")) for message in messages if message["role"] == "system"]
-    conversation = "\n".join(
-        f"{message['role']}: {text_content(message.get('content'))}"
-        for message in messages
-        if message["role"] != "system"
-    )
-    example = {
-        "input": conversation,
-        "output": "",
-        "input_assets": image_assets(messages),
-    }
-    with GENERATION_LOCK:
-        result = generate_multimodal_one(
-            MODEL,
-            PROCESSOR,
-            "\n\n".join(part for part in system_parts if part),
-            example,
-            generation,
-            Path.cwd(),
-            CHAT_TEMPLATE_KWARGS,
-        )
-    # The evaluator does not currently expose token counts for multimodal
-    # requests. Returning zero is preferable to inventing usage data.
-    return str(result["actual"]), 0, 0
-
-
 def bounded_number(value: Any, default: float, minimum: float, maximum: float) -> float:
     number = float(default if value is None else value)
     if not math.isfinite(number) or number < minimum or number > maximum:
@@ -220,8 +169,7 @@ class Handler(BaseHTTPRequestHandler):
         if not API_KEY:
             return True
         supplied = self.headers.get("authorization", "")
-        expected = "Bearer " + API_KEY
-        return hmac.compare_digest(supplied, expected)
+        return hmac.compare_digest(supplied, "Bearer " + API_KEY)
 
     def send_json(self, status: int, payload: Any) -> None:
         body = json.dumps(payload).encode("utf-8")
@@ -274,12 +222,7 @@ class Handler(BaseHTTPRequestHandler):
                 return
             started = time.perf_counter()
             try:
-                if MODEL_LOADER == "image_text_to_text":
-                    content, prompt_tokens, completion_tokens = generate_multimodal(messages, generation)
-                else:
-                    if image_assets(messages):
-                        raise ValueError("This text model does not accept image content parts.")
-                    content, prompt_tokens, completion_tokens = generate_text(messages, generation)
+                content, prompt_tokens, completion_tokens = generate_text(messages, generation)
             finally:
                 REQUEST_SLOTS.release()
             latency_ms = round((time.perf_counter() - started) * 1000)

package/training/local-runner/src/sft_data.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+
+IGNORE_INDEX = -100
+
+
+def _token_ids(value: Any, description: str) -> list[int]:
+    if isinstance(value, Mapping):
+        value = value.get("input_ids")
+    if not isinstance(value, Sequence) or isinstance(value, (str, bytes)):
+        raise TypeError(f"{description} did not return a token-id sequence")
+    token_ids = list(value)
+    if not token_ids or any(not isinstance(token_id, int) for token_id in token_ids):
+        raise ValueError(f"{description} returned empty or invalid token ids")
+    return token_ids
+
+
+def build_assistant_only_example(
+    tokenizer: Any,
+    messages: Any,
+    *,
+    max_length: int,
+) -> dict[str, list[int]]:
+    """Tokenize one chat row while computing loss only on the final assistant turn.
+
+    The entire assistant completion is reserved before the prompt is truncated.
+    Prompt truncation therefore removes the oldest prompt tokens and can never
+    silently discard the expected answer.
+    """
+    if max_length < 2:
+        raise ValueError("max_length must be at least 2")
+    if not isinstance(messages, list) or len(messages) < 2:
+        raise ValueError("messages must contain prompt context and a final assistant turn")
+    if any(not isinstance(message, dict) for message in messages):
+        raise ValueError("every message must be an object")
+    allowed_roles = {"system", "user", "assistant", "tool"}
+    for message in messages:
+        role = message.get("role")
+        if role not in allowed_roles:
+            raise ValueError(f"unsupported message role: {role!r}")
+        if not isinstance(message.get("content"), str):
+            raise ValueError("the bundled trainer accepts text-only message content")
+
+    assistant = messages[-1]
+    if assistant.get("role") != "assistant":
+        raise ValueError("the final message must have role=assistant")
+    content = assistant.get("content")
+    if not isinstance(content, str) or not content.strip():
+        raise ValueError("the final assistant message must contain non-empty text")
+
+    try:
+        prompt_ids = _token_ids(
+            tokenizer.apply_chat_template(
+                messages[:-1],
+                tokenize=True,
+                add_generation_prompt=True,
+                return_dict=False,
+            ),
+            "prompt chat template",
+        )
+        full_ids = _token_ids(
+            tokenizer.apply_chat_template(
+                messages,
+                tokenize=True,
+                add_generation_prompt=False,
+                return_dict=False,
+            ),
+            "full chat template",
+        )
+    except Exception as exc:
+        raise ValueError(f"Qwen chat-template rendering failed: {exc}") from exc
+
+    if full_ids[: len(prompt_ids)] != prompt_ids:
+        raise ValueError(
+            "Qwen chat-template output is not prefix-aligned; refusing to guess the assistant loss boundary"
+        )
+    completion_ids = full_ids[len(prompt_ids) :]
+    if not completion_ids:
+        raise ValueError("Qwen chat template produced no assistant completion tokens")
+    if len(completion_ids) >= max_length:
+        raise ValueError(
+            f"assistant completion requires {len(completion_ids)} tokens, but max_length={max_length}; "
+            "increase max_seq_length so the answer is not truncated"
+        )
+
+    prompt_budget = max_length - len(completion_ids)
+    retained_prompt_ids = prompt_ids[-prompt_budget:]
+    input_ids = [*retained_prompt_ids, *completion_ids]
+    labels = [IGNORE_INDEX] * len(retained_prompt_ids) + completion_ids.copy()
+    return {
+        "input_ids": input_ids,
+        "attention_mask": [1] * len(input_ids),
+        "labels": labels,
+    }