Hugging Face's logo

KerzgrAI
/
PromptTercet-0.5B

Text Generation
Safetensors
PyTorch
English
tiny_gdn
senryu2
causal-lm
prompter
user-simulator
synthetic-data
chat
hybrid
gated-deltanet-2
parallax-attention
n-gram-memory
senryu-2
conversational
Model card Files
xet
 Community
PromptTercet-0.5B / inference.py
kerzgrr's picture
kerzgrr
Upload PromptTercet-0.5B (EMA weights)
f2e0f25 verified
Raw History Blame Contribute Delete
40.6 kB
#!/usr/bin/env python3
"""Standalone inference for KerzgrAI/PromptTercet-0.5B.
PromptTercet writes the user's side of a chat: give it a conversation and it
asks the next question, give it nothing and it opens one. You, or any
OpenAI-compatible model, write the assistant's side.
Examples:
python inference.py # it asks, you answer
python inference.py --system "You are a cooking assistant."
python inference.py --conversation chat.json # print the next question
python inference.py --count 500 --output prompts.jsonl
python inference.py --count 100 --turns 3 --output chats.jsonl \\
--answer-url https://openrouter.ai/api/v1 --answer-model openai/gpt-4o-mini
"""
from __future__ import annotations
import argparse
import importlib
import json
import os
import platform
import secrets
import shutil
import subprocess
import sys
import threading
import time
import unicodedata
import urllib.error
import urllib.request
import warnings
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path
from types import SimpleNamespace
from typing import Any, TextIO
import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from tokenizers import Tokenizer
def _silence_runtime_warnings() -> None:
patterns = (
r"tl\.make_block_ptr is deprecated",
r"Memory efficient kernel not used because",
r"Memory Efficient attention has been runtime disabled",
r"Flash attention kernel not used because",
r"Torch was not compiled with flash attention",
r"cuDNN attention kernel not used because",
r"cuDNN attention has been runtime disabled",
)
for pattern in patterns:
warnings.filterwarnings("ignore", message=pattern)
REPO_ID = "KerzgrAI/PromptTercet-0.5B"
MODEL_NAME = "PromptTercet-0.5B"
FLA_COMMIT = "cbb0a72efb55c18ca0ef4f298298317573ad2cb3"
FLA_REPO = "https://github.com/fla-org/flash-linear-attention.git"
FLA_RUNTIME_REQUIREMENT = "transformers>=4.45.0"
# Parallax attention is not in a flash-linear-attention release yet, so fla
# installs from git at the commit the model was trained with.
FLA_MODULES = (
"fla.layers.gdn2",
"fla.models.utils",
"fla.modules.rotary",
"fla.ops.parallax.parallel",
"fla.ops.parallax.decode",
)
PATCH_FILES = (
"fla/__init__.py",
"fla/ops/__init__.py",
"fla/layers/__init__.py",
"fla/ops/simple_gla/__init__.py",
)
TINY_GDN_FILES = (
"tiny_gdn/__init__.py",
"tiny_gdn/config.py",
"tiny_gdn/model.py",
"tiny_gdn/haiku_layers.py",
"tiny_gdn/nn_common.py",
"tiny_gdn/fused_ops.py",
"tiny_gdn/senryu2.py",
"tiny_gdn/parallax_ops.py",
"tiny_gdn/gated_residual_ops.py",
"tiny_gdn/ngram.py",
"tiny_gdn/tools.py",
"tiny_gdn/smoltalk_chat.py",
"tiny_gdn/chatml.py",
"tiny_gdn/detokenize.py",
"tiny_gdn/repetition.py",
)
NGRAM_TABLES_FILE = "ngram_tables.safetensors"
ANSWER_REQUEST_RETRIES = 6
# A conversation's own turns, as PromptTercet's chat template expects them:
# its questions are the template's assistant turns and the replies its user turns.
PROMPTER_ROLES = {"system": "system", "user": "assistant", "assistant": "user"}
@dataclass(frozen=True)
class Source:
repo_id: str
local_dir: Path | None
# A full copy of this repo next to inference.py, used instead of the Hub.
local_copy: Path | None
def _local_copy() -> Path | None:
here = Path(__file__).resolve().parent
config = here / "config.json"
if not all((here / name).is_file() for name in ("config.json", "model.safetensors", NGRAM_TABLES_FILE)):
return None
try:
name = json.loads(config.read_text(encoding="utf-8")).get("model_name")
except (OSError, ValueError):
return None
return here if name == MODEL_NAME else None
def _download(source: Source, filename: str) -> Path:
if source.local_copy is not None and (source.local_copy / filename).is_file():
return source.local_copy / filename
return Path(
hf_hub_download(
repo_id=source.repo_id,
filename=filename,
local_dir=str(source.local_dir) if source.local_dir else None,
)
)
def _run(cmd: list[str], *, cwd: Path | None = None) -> None:
print("+", " ".join(cmd), flush=True)
merged = os.environ.copy()
merged.setdefault("PYTHONUTF8", "1")
merged.setdefault("PYTHONIOENCODING", "utf-8")
subprocess.check_call(cmd, cwd=str(cwd) if cwd else None, env=merged)
def _pip_install(*args: str) -> None:
_run([sys.executable, "-m", "pip", "install", *args])
def _fla_importable() -> tuple[bool, str]:
try:
for name in FLA_MODULES:
importlib.import_module(name)
except Exception as error: # noqa: BLE001
return False, str(error)
return True, ""
def _cache_root() -> Path:
override = os.environ.get("PROMPTTERCET_CACHE")
path = Path(override).expanduser().resolve() if override else Path.home() / ".cache" / "prompttercet"
path.mkdir(parents=True, exist_ok=True)
return path
def _install_fla(source: Source) -> None:
print("flash-linear-attention missing/broken — installing automatically …", flush=True)
_pip_install("einops", "numpy", FLA_RUNTIME_REQUIREMENT)
if platform.system() != "Windows":
_pip_install("--no-deps", "--force-reinstall", f"git+{FLA_REPO}@{FLA_COMMIT}")
return
if shutil.which("git") is None:
raise RuntimeError("git is required to auto-install flash-linear-attention. Install Git and put it on PATH.")
fla_root = _cache_root() / "flash-linear-attention"
if not (fla_root / ".git").is_dir():
if fla_root.exists():
shutil.rmtree(fla_root)
_run(["git", "clone", "--filter=blob:none", FLA_REPO, str(fla_root)])
_run(["git", "fetch", "--depth", "1", "origin", FLA_COMMIT], cwd=fla_root)
_run(["git", "checkout", "--force", FLA_COMMIT], cwd=fla_root)
print("Applying Windows FLA import patches …", flush=True)
for relative in PATCH_FILES:
target = fla_root / relative
target.parent.mkdir(parents=True, exist_ok=True)
shutil.copy2(_download(source, f"windows_fla_patches/{relative}"), target)
_pip_install("--no-build-isolation", "--no-deps", "-e", str(fla_root))
def _ensure_fla(source: Source) -> None:
ok, error = _fla_importable()
if ok:
return
print(f"FLA not ready ({error})", flush=True)
try:
_install_fla(source)
except Exception as install_error:
raise RuntimeError(
"Automatic flash-linear-attention install failed.\n"
f"Original import error: {error}\n"
f"Install error: {install_error}"
) from install_error
for name in list(sys.modules):
if name == "fla" or name.startswith("fla."):
del sys.modules[name]
importlib.invalidate_caches()
ok, error = _fla_importable()
if not ok:
raise RuntimeError(f"flash-linear-attention installed but still failed to import: {error}")
print("flash-linear-attention ready.", flush=True)
def _ensure_tiny_gdn(source: Source) -> Path:
for root in (Path(__file__).resolve().parent, source.local_copy, source.local_dir):
if root is not None and all((root / name).is_file() for name in TINY_GDN_FILES):
return root
for name in TINY_GDN_FILES:
_download(source, name)
return _download(source, "tiny_gdn/__init__.py").parent.parent
def _import_stack() -> SimpleNamespace:
# tiny_gdn is fetched from the Hub and FLA may be installed at runtime, so
# these modules only become importable after _ensure_tiny_gdn/_ensure_fla.
config = importlib.import_module("tiny_gdn.config")
senryu2 = importlib.import_module("tiny_gdn.senryu2")
chatml = importlib.import_module("tiny_gdn.chatml")
detokenize = importlib.import_module("tiny_gdn.detokenize")
repetition = importlib.import_module("tiny_gdn.repetition")
return SimpleNamespace(
TinyGDNConfig=config.TinyGDNConfig,
Senryu2ForCausalLM=senryu2.Senryu2ForCausalLM,
encode_prompter_messages=chatml.encode_prompter_messages,
IncrementalUtf8Decoder=detokenize.IncrementalUtf8Decoder,
apply_repetition_penalty=repetition.apply_repetition_penalty,
structural_token_flags=repetition.structural_token_flags,
)
@dataclass(frozen=True)
class Sampling:
temperature: float
top_p: float
top_k: int
repetition_penalty: float
repetition_window: int
max_new_tokens: int
context_length: int
@dataclass(frozen=True)
class Question:
text: str
tokens: int
# "stop" ended the turn itself; "length" hit --max-new-tokens and is cut off.
stop_reason: str
@property
def usable(self) -> bool:
return self.stop_reason == "stop" and bool(self.text)
def _sample(logits: torch.Tensor, sampling: Sampling, generator: torch.Generator) -> int:
logits = logits.float()
if sampling.temperature <= 1e-5:
return int(torch.argmax(logits).item())
logits = logits / sampling.temperature
if 0 < sampling.top_k < logits.shape[-1]:
threshold = torch.topk(logits, sampling.top_k).values[-1]
logits = logits.masked_fill(logits < threshold, -torch.inf)
if sampling.top_p < 1.0:
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
remove = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > sampling.top_p
remove[1:] = remove[:-1].clone()
remove[0] = False
sorted_logits = sorted_logits.masked_fill(remove, -torch.inf)
logits = torch.full_like(logits, -torch.inf)
logits.scatter_(0, sorted_indices, sorted_logits)
probs = torch.softmax(logits, dim=-1).cpu()
return int(torch.multinomial(probs, 1, generator=generator).item())
def prompter_messages(conversation: list[dict[str, str]]) -> list[dict[str, str]]:
return [{"role": PROMPTER_ROLES[item["role"]], "content": item["content"]} for item in conversation]
class Prompter:
"""PromptTercet on one GPU. Generations run one at a time, so threads can share it."""
def __init__(self, model: Any, tokenizer: Tokenizer, stack: SimpleNamespace, device: torch.device, sampling: Sampling, seed: int | None) -> None:
self.model = model
self.tokenizer = tokenizer
self.stack = stack
self.device = device
self.sampling = sampling
self.stop_ids = {int(model.config.eos_token_id)}
im_end = tokenizer.token_to_id("<|im_end|>")
if im_end is not None:
self.stop_ids.add(im_end)
# The penalty skips whitespace, punctuation and chat tokens, so turns can still end.
self.structural = stack.structural_token_flags(tokenizer) if sampling.repetition_penalty != 1.0 else ()
self._seed = seed
self._calls = 0
self._lock = threading.Lock()
def prompt_length(self, conversation: list[dict[str, str]]) -> int:
return len(self.stack.encode_prompter_messages(self.tokenizer, prompter_messages(conversation)))
def ask(
self,
conversation: list[dict[str, str]],
*,
on_text: Callable[[str], None] | None = None,
stop: threading.Event | None = None,
) -> Question:
prompt_ids = self.stack.encode_prompter_messages(self.tokenizer, prompter_messages(conversation))
with self._lock:
if self._seed is None:
seed = secrets.randbits(63)
else:
seed = self._seed + self._calls
self._calls += 1
return self._generate(prompt_ids, seed, on_text, stop)
@torch.inference_mode()
def _generate(
self,
prompt_ids: list[int],
seed: int,
on_text: Callable[[str], None] | None,
stop: threading.Event | None,
) -> Question:
sampling = self.sampling
token_ids = list(prompt_ids[-sampling.context_length :])
generated = 0
decoder = self.stack.IncrementalUtf8Decoder(self.tokenizer)
generator = torch.Generator(device="cpu")
generator.manual_seed(seed)
stop_reason = "length"
try:
output = self.model(
torch.tensor([token_ids], dtype=torch.long, device=self.device),
return_logits=True,
logits_to_keep=1,
use_cache=True,
)
for step in range(sampling.max_new_tokens):
if stop is not None and stop.is_set():
stop_reason = "interrupted"
break
if output.logits is None:
raise RuntimeError("Model returned no logits")
logits = output.logits[0, -1].float()
if self.structural:
self.stack.apply_repetition_penalty(
logits,
token_ids,
penalty=sampling.repetition_penalty,
window=sampling.repetition_window,
structural=self.structural,
)
next_id = _sample(logits, sampling, generator)
if next_id in self.stop_ids:
stop_reason = "stop"
break
token_ids.append(next_id)
generated += 1
delta = decoder.push(next_id)
if delta and on_text is not None:
on_text(delta)
if step + 1 == sampling.max_new_tokens:
break
if output.past_key_values is None:
raise RuntimeError("Model returned no decode cache")
if output.past_key_values.length + 1 > self.model.config.max_position_embeddings:
stop_reason = "context_full"
break
output = self.model(
torch.tensor([[next_id]], dtype=torch.long, device=self.device),
return_logits=True,
logits_to_keep=1,
past_key_values=output.past_key_values,
use_cache=True,
)
except KeyboardInterrupt:
stop_reason = "interrupted"
text = decoder.finalize().strip()
return Question(text=text, tokens=generated, stop_reason=stop_reason)
class AnswerError(RuntimeError):
pass
@dataclass(frozen=True)
class AnswerModel:
"""An OpenAI-compatible chat completions endpoint that writes the assistant turns."""
base_url: str
model: str
api_key: str
max_tokens: int
temperature: float
extra: dict[str, Any]
timeout: float
def reply(self, messages: list[dict[str, str]], stop: threading.Event) -> tuple[str, str]:
"""``(content, finish_reason)`` of one completion, retrying transient failures."""
body = {
**self.extra,
"model": self.model,
"messages": messages,
"max_tokens": self.max_tokens,
"temperature": self.temperature,
}
headers = {"Content-Type": "application/json", "User-Agent": "PromptTercet-0.5B/inference.py"}
if self.api_key:
headers["Authorization"] = f"Bearer {self.api_key}"
url = f"{self.base_url.rstrip('/')}/chat/completions"
last = "no attempt"
for attempt in range(ANSWER_REQUEST_RETRIES):
if stop.is_set():
raise AnswerError("stopped")
delay = min(2.0**attempt, 30.0)
try:
request = urllib.request.Request(url, data=json.dumps(body).encode("utf-8"), headers=headers, method="POST")
with urllib.request.urlopen(request, timeout=self.timeout) as response:
payload = json.loads(response.read().decode("utf-8"))
except urllib.error.HTTPError as error:
detail = error.read().decode("utf-8", errors="replace")[:300]
last = f"HTTP {error.code}: {detail}"
if error.code not in {408, 409, 425, 429} and error.code < 500:
raise AnswerError(last) from error
retry_after = error.headers.get("Retry-After") if error.headers else None
if retry_after and retry_after.isdigit():
delay = min(float(retry_after), 60.0)
except (urllib.error.URLError, TimeoutError, ConnectionError, OSError, ValueError) as error:
last = f"{type(error).__name__}: {error}"
else:
if isinstance(payload, dict) and payload.get("error"):
last = f"API error: {json.dumps(payload['error'])[:300]}"
else:
try:
choice = payload["choices"][0]
content = choice["message"].get("content") or ""
except (KeyError, IndexError, TypeError, AttributeError) as error:
raise AnswerError(f"Unexpected response: {json.dumps(payload)[:300]}") from error
return str(content).strip(), str(choice.get("finish_reason") or "")
if stop.wait(delay):
raise AnswerError("stopped")
raise AnswerError(f"{ANSWER_REQUEST_RETRIES} requests failed; last: {last}")
def opening_key(text: str) -> str:
return unicodedata.normalize("NFC", text).strip()
def load_conversation(path: str) -> list[dict[str, str]]:
raw = sys.stdin.read() if path == "-" else Path(path).read_text(encoding="utf-8")
payload = json.loads(raw)
messages = payload.get("messages") if isinstance(payload, dict) else payload
if not isinstance(messages, list):
raise ValueError("A conversation is a JSON list of messages or an object with a 'messages' list")
conversation: list[dict[str, str]] = []
for index, item in enumerate(messages):
if not isinstance(item, dict) or item.get("role") not in PROMPTER_ROLES:
raise ValueError(f"Message {index} needs a role of system, user or assistant")
content = item.get("content")
if not isinstance(content, str) or not content.strip():
raise ValueError(f"Message {index} needs non-empty text content")
if item["role"] == "system" and index != 0:
raise ValueError("Only the first message can be a system prompt")
conversation.append({"role": item["role"], "content": content})
if conversation and conversation[-1]["role"] == "user":
raise ValueError("The conversation ends with a user turn. PromptTercet writes the user turn after an assistant reply.")
return conversation
class Style:
def __init__(self, enabled: bool) -> None:
self.enabled = enabled
def wrap(self, code: str, text: str) -> str:
return f"\033[{code}m{text}\033[0m" if self.enabled else text
def user(self, text: str) -> str:
return self.wrap("1;36", text)
def assistant(self, text: str) -> str:
return self.wrap("1;32", text)
def note(self, text: str) -> str:
return self.wrap("2", text)
INTERACTIVE_HELP = (
"Commands: /retry asks again, /system TEXT sets the system prompt and starts over, "
"/reset starts over, /save FILE appends the chat as a JSON line, /exit quits."
)
def run_interactive(prompter: Prompter, system: str, style: Style) -> int:
print(
f"{MODEL_NAME} plays the user and you play the assistant: answer each question it asks.\n"
f"{INTERACTIVE_HELP} Ctrl+C stops a question.",
flush=True,
)
def fresh(prompt: str) -> list[dict[str, str]]:
return [{"role": "system", "content": prompt}] if prompt.strip() else []
conversation = fresh(system)
ask_now = True
while True:
if ask_now:
ask_now = False
if prompter.prompt_length(conversation) > prompter.sampling.context_length:
print(style.note(f"(the chat is over {prompter.sampling.context_length} tokens; the oldest are dropped)"))
print(f"\n{style.user('user>')} ", end="", flush=True)
try:
question = prompter.ask(conversation, on_text=lambda delta: print(delta, end="", flush=True))
except KeyboardInterrupt:
question = Question(text="", tokens=0, stop_reason="interrupted")
print(flush=True)
if question.stop_reason == "interrupted":
print(style.note("(stopped; /retry asks again)"))
elif not question.text:
print(style.note("(empty question; /retry asks again)"))
else:
if question.stop_reason == "length":
print(style.note(f"(cut off at --max-new-tokens {prompter.sampling.max_new_tokens}; /retry asks again)"))
conversation.append({"role": "user", "content": question.text})
try:
line = input(f"{style.assistant('assistant>')} ")
except (EOFError, KeyboardInterrupt):
print()
return 0
text = line.strip()
if not text:
continue
if text.startswith("/"):
command, _, rest = text.partition(" ")
if command in {"/exit", "/quit"}:
return 0
if command == "/help":
print(INTERACTIVE_HELP)
elif command == "/reset":
conversation = fresh(system)
ask_now = True
print(style.note("(new chat)"))
elif command == "/system":
system = rest.strip()
conversation = fresh(system)
ask_now = True
print(style.note("(system prompt set; new chat)" if system else "(system prompt cleared; new chat)"))
elif command == "/retry":
if conversation and conversation[-1]["role"] == "user":
conversation.pop()
ask_now = True
elif command == "/save":
path = Path(rest.strip() or "conversations.jsonl")
if not any(item["role"] == "user" for item in conversation):
print(style.note("(nothing to save yet)"))
continue
with path.open("a", encoding="utf-8") as handle:
handle.write(json.dumps({"messages": conversation}, ensure_ascii=False) + "\n")
print(style.note(f"(appended to {path})"))
else:
print(style.note(f"(unknown command {command}; /help lists them)"))
continue
if not conversation or conversation[-1]["role"] != "user":
print(style.note("(no question to answer; /retry asks again)"))
continue
conversation.append({"role": "assistant", "content": text})
ask_now = True
def run_next_question(prompter: Prompter, conversation: list[dict[str, str]]) -> int:
question = prompter.ask(conversation)
if question.stop_reason == "interrupted":
return 130
print(question.text, flush=True)
if not question.usable:
print(f"[warn] question ended with {question.stop_reason}", file=sys.stderr, flush=True)
return 1
return 0
@dataclass
class BatchState:
count: int
written: int = 0
in_flight: int = 0
failed: int = 0
redone_questions: int = 0
repeated_openings: int = 0
redone_answers: int = 0
last_error: str = ""
seen: set[str] = field(default_factory=set)
lock: threading.Lock = field(default_factory=threading.Lock)
def claim(self) -> bool:
with self.lock:
if self.written + self.in_flight >= self.count:
return False
self.in_flight += 1
return True
def bump(self, name: str) -> None:
with self.lock:
setattr(self, name, getattr(self, name) + 1)
def note(self, error: str) -> None:
with self.lock:
self.last_error = error
def finish(self, ok: bool, stopping: bool) -> bool:
"""Settle one claimed chat; True when failures say the batch should stop."""
with self.lock:
self.in_flight -= 1
if ok:
self.written += 1
elif not stopping:
self.failed += 1
return self.failed >= 20 and self.failed > self.written
def summary(self) -> str:
parts = [f"{self.written}/{self.count} written"]
if self.failed:
parts.append(f"{self.failed} failed")
if self.repeated_openings:
parts.append(f"{self.repeated_openings} repeated openings resampled")
if self.redone_questions:
parts.append(f"{self.redone_questions} cut-off questions resampled")
if self.redone_answers:
parts.append(f"{self.redone_answers} cut-off answers redone")
return ", ".join(parts)
def seen_openings(path: Path) -> set[str]:
seen: set[str] = set()
if not path.is_file():
return seen
with path.open(encoding="utf-8") as handle:
for line in handle:
try:
messages = json.loads(line).get("messages", [])
except (ValueError, AttributeError):
continue
first = next((item for item in messages if isinstance(item, dict) and item.get("role") == "user"), None)
if first and isinstance(first.get("content"), str):
seen.add(opening_key(first["content"]))
return seen
def run_batch(prompter: Prompter, args: argparse.Namespace, answer: AnswerModel | None) -> int:
output_path = Path(args.output) if args.output else None
state = BatchState(count=args.count)
if output_path is not None:
state.seen = seen_openings(output_path)
if state.seen:
print(f"{output_path} already has {len(state.seen)} openings; new ones will differ.", file=sys.stderr)
sink: TextIO = output_path.open("a", encoding="utf-8") if output_path else sys.stdout
write_lock = threading.Lock()
stop = threading.Event()
system = args.system.strip()
answer_system = system if args.answer_system is None else args.answer_system.strip()
def ask(conversation: list[dict[str, str]], opening: bool) -> str | None:
for _ in range(args.attempts):
question = prompter.ask(conversation, stop=stop)
if question.stop_reason == "interrupted" or stop.is_set():
return None
if not question.usable:
state.bump("redone_questions")
continue
if opening:
key = opening_key(question.text)
with state.lock:
if key in state.seen:
state.repeated_openings += 1
continue
state.seen.add(key)
return question.text
state.note(f"no usable question in {args.attempts} tries")
return None
def reply(conversation: list[dict[str, str]]) -> str | None:
assert answer is not None
messages = ([{"role": "system", "content": answer_system}] if answer_system else []) + [
item for item in conversation if item["role"] != "system"
]
for _ in range(args.attempts):
try:
content, finish = answer.reply(messages, stop)
except AnswerError as error:
if not stop.is_set():
state.note(f"answer model: {error}")
return None
if content and finish != "length":
return content
state.bump("redone_answers")
state.note(f"{args.attempts} answers in a row hit --answer-max-tokens {args.answer_max_tokens} or came back empty")
return None
def conversation_once() -> bool:
"""One chat of exactly --turns questions, written only when complete."""
conversation = [{"role": "system", "content": system}] if system else []
for turn in range(args.turns):
question = ask(conversation, opening=turn == 0)
if question is None:
return False
conversation.append({"role": "user", "content": question})
if answer is not None:
text = reply(conversation)
if text is None:
return False
conversation.append({"role": "assistant", "content": text})
if stop.is_set():
return False
with write_lock:
sink.write(json.dumps({"messages": conversation}, ensure_ascii=False) + "\n")
sink.flush()
return True
def worker() -> None:
while not stop.is_set() and state.claim():
ok = False
try:
ok = conversation_once()
except Exception as error: # noqa: BLE001
state.note(f"{type(error).__name__}: {error}")
if state.finish(ok, stop.is_set()):
state.note(f"stopped after {state.failed} failed chats; last: {state.last_error}")
stop.set()
workers = [threading.Thread(target=worker, daemon=True) for _ in range(max(1, args.workers))]
for thread in workers:
thread.start()
tty = sys.stderr.isatty()
shown = ""
shown_at = 0.0
interrupted = False
try:
while any(thread.is_alive() for thread in workers):
for thread in workers:
thread.join(timeout=0.5)
line = state.summary()
if line != shown and (tty or time.monotonic() - shown_at >= 15):
shown, shown_at = line, time.monotonic()
print(f"\r[{MODEL_NAME}] {line}\033[K" if tty else f"[{MODEL_NAME}] {line}", end="" if tty else "\n", file=sys.stderr, flush=True)
if tty:
print(file=sys.stderr)
except KeyboardInterrupt:
interrupted = True
stop.set()
print("\nStopping after the requests in flight (Ctrl+C again quits now) …", file=sys.stderr, flush=True)
try:
for thread in workers:
thread.join()
except KeyboardInterrupt:
pass
finally:
if output_path is not None:
sink.close()
print(f"[{MODEL_NAME}] done: {state.summary()}", file=sys.stderr, flush=True)
if state.last_error and state.written < state.count:
print(f"[{MODEL_NAME}] last failure: {state.last_error}", file=sys.stderr, flush=True)
if interrupted:
return 130
return 0 if state.written >= state.count else 1
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description=f"{MODEL_NAME}: writes the user's side of a chat",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
mode = parser.add_argument_group("what to do (default: an interactive chat where you answer)")
mode.add_argument("--system", default="", help="System prompt of the chat; it steers what PromptTercet asks about")
mode.add_argument("--conversation", default=None, help="JSON chat (list of messages, or {messages: [...]}); print the next question. '-' reads stdin")
mode.add_argument("--count", type=int, default=0, help="Batch mode: write this many chats as JSON lines")
mode.add_argument("--turns", type=int, default=1, help="Questions per chat in batch mode; above 1 needs --answer-url")
mode.add_argument("--output", default=None, help="JSONL file for batch mode (default: stdout). Appends, and new openings differ from those already in it")
mode.add_argument("--attempts", type=int, default=8, help="Samples per turn before a chat is given up: cut-off or repeated questions, cut-off answers")
mode.add_argument("--workers", type=int, default=None, help="Chats in flight in batch mode (default: 8 with an answer model, else 1)")
answers = parser.add_argument_group("answer model for batch mode (any OpenAI-compatible API)")
answers.add_argument("--answer-url", default=None, help="Base URL, e.g. https://openrouter.ai/api/v1 or http://localhost:8000/v1")
answers.add_argument("--answer-model", default=None, help="Model name sent to the API")
answers.add_argument("--answer-key-env", default="OPENAI_API_KEY", help="Environment variable holding the API key (unset sends none)")
answers.add_argument("--answer-system", default=None, help="System prompt only the answer model sees, not saved (default: --system)")
answers.add_argument("--answer-max-tokens", type=int, default=4096, help="Answers that hit this are discarded and redone")
answers.add_argument("--answer-temperature", type=float, default=0.7)
answers.add_argument("--answer-extra", default="{}", help='Extra request fields as JSON, e.g. \'{"reasoning": {"effort": "low"}}\'')
answers.add_argument("--answer-timeout", type=float, default=600.0, help="Seconds per request")
sampling = parser.add_argument_group("PromptTercet sampling")
sampling.add_argument("--temperature", type=float, default=0.6)
sampling.add_argument("--top-p", type=float, default=0.95)
sampling.add_argument("--top-k", type=int, default=0, help="0 turns top-k off")
sampling.add_argument(
"--repetition-penalty",
type=float,
default=1.36,
help="Penalty on words already used twice in the chat, which keeps follow-ups from repeating it (1 turns it off)",
)
sampling.add_argument("--repetition-window", type=int, default=0, help="Tokens the penalty looks back over (0: the whole chat)")
sampling.add_argument("--max-new-tokens", type=int, default=1024, help="Questions that hit this are cut off")
sampling.add_argument("--context-length", type=int, default=8192, help="Chat tokens kept, oldest dropped first (up to 32768)")
sampling.add_argument("--seed", type=int, default=None, help="Fix the sampling seed (default: a fresh seed per question)")
runtime = parser.add_argument_group("runtime")
runtime.add_argument("--device", default="cuda", help=f"CUDA device (cuda, cuda:1, …); {MODEL_NAME} needs a GPU")
runtime.add_argument("--ngram-device", default=None, help="Where the 2.2 GB n-gram tables live: the GPU (default, --device) or cpu")
runtime.add_argument("--no-color", action="store_true")
runtime.add_argument("--local-dir", default=None, help="Download the model files into this folder")
runtime.add_argument("--repo-id", default=REPO_ID)
return parser.parse_args()
def validate(args: argparse.Namespace) -> str | None:
if args.count < 0 or args.turns < 1 or args.attempts < 1:
return "--count must be 0 or more, --turns and --attempts at least 1"
if args.workers is not None and args.workers < 1:
return "--workers must be at least 1"
if args.count and args.conversation:
return "Use --count or --conversation, not both"
if bool(args.answer_url) != bool(args.answer_model):
return "--answer-url and --answer-model go together"
if args.answer_url and not args.count:
return "--answer-url is for batch mode; add --count"
if args.answer_url and not args.answer_url.startswith(("http://", "https://")):
return "--answer-url must start with http:// or https://"
if args.answer_max_tokens < 1 or args.answer_timeout <= 0:
return "--answer-max-tokens and --answer-timeout must be positive"
if args.count and args.turns > 1 and not args.answer_url:
return "--turns above 1 needs an answer model: add --answer-url and --answer-model"
if not 1 <= args.max_new_tokens <= 32768 or not 256 <= args.context_length <= 32768:
return "--max-new-tokens must be 1-32768 and --context-length 256-32768"
if args.temperature < 0 or not 0 < args.top_p <= 1 or args.top_k < 0 or args.repetition_penalty <= 0:
return "Sampling needs --temperature >= 0, 0 < --top-p <= 1, --top-k >= 0 and --repetition-penalty > 0"
if args.repetition_window < 0:
return "--repetition-window must be 0 (the whole chat) or more"
try:
extra = json.loads(args.answer_extra)
except json.JSONDecodeError as error:
return f"--answer-extra is not JSON: {error}"
if not isinstance(extra, dict):
return "--answer-extra must be a JSON object"
return None
def main() -> int:
_silence_runtime_warnings()
args = parse_args()
problem = validate(args)
if problem:
print(problem, file=sys.stderr)
return 2
conversation: list[dict[str, str]] | None = None
if args.conversation is not None:
try:
conversation = load_conversation(args.conversation)
except (OSError, ValueError) as error:
print(f"Cannot read --conversation: {error}", file=sys.stderr)
return 2
if args.system.strip() and not (conversation and conversation[0]["role"] == "system"):
conversation.insert(0, {"role": "system", "content": args.system.strip()})
try:
device = torch.device(args.device)
ngram_device = torch.device(args.ngram_device) if args.ngram_device else device
except RuntimeError as error:
print(f"Invalid device: {error}", file=sys.stderr)
return 2
if device.type != "cuda" or not torch.cuda.is_available():
print(
f"{MODEL_NAME} needs a CUDA GPU: its Gated DeltaNet-2 and Parallax layers run on "
f"flash-linear-attention Triton kernels. torch {torch.__version__} reports "
f"CUDA available={torch.cuda.is_available()}.",
file=sys.stderr,
)
return 1
if ngram_device.type not in {"cuda", "cpu"}:
print("--ngram-device must be a CUDA device or cpu", file=sys.stderr)
return 2
source = Source(
repo_id=args.repo_id,
local_dir=Path(args.local_dir).resolve() if args.local_dir else None,
local_copy=_local_copy() if args.repo_id == REPO_ID else None,
)
where = source.local_copy or f"huggingface.co/{source.repo_id}"
print(f"Loading {MODEL_NAME} from {where} …", file=sys.stderr, flush=True)
try:
package_root = _ensure_tiny_gdn(source)
if str(package_root) not in sys.path:
sys.path.insert(0, str(package_root))
_ensure_fla(source)
stack = _import_stack()
except Exception as error: # noqa: BLE001
print(f"Setup failed: {error}", file=sys.stderr)
return 1
config = stack.TinyGDNConfig.from_json(_download(source, "config.json"))
with torch.device(device):
model = stack.Senryu2ForCausalLM(config)
model = model.to(dtype=torch.bfloat16)
model.load_state_dict(load_file(str(_download(source, "model.safetensors")), device=str(device)), strict=True)
if model.ngram is not None:
model.ngram.load_tables(_download(source, NGRAM_TABLES_FILE), device=ngram_device, dtype=torch.bfloat16)
model.eval()
model.requires_grad_(False)
tokenizer = Tokenizer.from_file(str(_download(source, "tokenizer.json")))
sampling = Sampling(
temperature=args.temperature,
top_p=args.top_p,
top_k=args.top_k,
repetition_penalty=args.repetition_penalty,
repetition_window=args.repetition_window,
max_new_tokens=args.max_new_tokens,
context_length=min(args.context_length, config.max_position_embeddings),
)
prompter = Prompter(model, tokenizer, stack, device, sampling, args.seed)
print(
f"Ready: {config.num_hidden_layers} layers on {device}, n-gram tables on {ngram_device}.",
file=sys.stderr,
flush=True,
)
if conversation is not None:
return run_next_question(prompter, conversation)
if args.count:
answer = None
if args.answer_url:
answer = AnswerModel(
base_url=args.answer_url,
model=args.answer_model,
api_key=os.environ.get(args.answer_key_env, "").strip(),
max_tokens=args.answer_max_tokens,
temperature=args.answer_temperature,
extra=json.loads(args.answer_extra),
timeout=args.answer_timeout,
)
if args.workers is None:
args.workers = 8 if answer else 1
return run_batch(prompter, args, answer)
style = Style(not args.no_color and sys.stdout.isatty())
return run_interactive(prompter, args.system, style)
if __name__ == "__main__":
raise SystemExit(main())