mirror of https://github.com/razor-ai/soup.git
74 lines
2.9 KiB
Python
74 lines
2.9 KiB
Python
"""Speculative decoding auto-pairing (v0.30.0).
|
|
|
|
Maps a target model to a known-good draft model. Draft+target overhead is
|
|
only worth it at ~30B+ — below that, pair returns None.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from typing import Optional
|
|
|
|
# Draft models picked for vocab / tokenizer match with the target family.
|
|
# Rule of thumb: draft 10-50x smaller than target, same architecture family.
|
|
_DRAFT_PAIRS: dict[str, str] = {
|
|
# Llama 3 family (share Llama-3 tokenizer)
|
|
"meta-llama/llama-3.1-70b": "meta-llama/Llama-3.2-1B",
|
|
"meta-llama/llama-3.1-70b-instruct": "meta-llama/Llama-3.2-1B-Instruct",
|
|
"meta-llama/llama-3.3-70b": "meta-llama/Llama-3.2-1B",
|
|
"meta-llama/llama-3.3-70b-instruct": "meta-llama/Llama-3.2-1B-Instruct",
|
|
"meta-llama/llama-3.1-405b": "meta-llama/Llama-3.2-3B",
|
|
"meta-llama/llama-3.1-405b-instruct": "meta-llama/Llama-3.2-3B-Instruct",
|
|
# Llama 4
|
|
"meta-llama/llama-4-scout-17b-16e": "meta-llama/Llama-3.2-1B-Instruct",
|
|
"meta-llama/llama-4-scout-17b-16e-instruct": "meta-llama/Llama-3.2-1B-Instruct",
|
|
"meta-llama/llama-4-maverick-17b-128e": "meta-llama/Llama-3.2-1B-Instruct",
|
|
# Qwen 2.5 / 3
|
|
"qwen/qwen2.5-72b": "Qwen/Qwen2.5-0.5B",
|
|
"qwen/qwen2.5-72b-instruct": "Qwen/Qwen2.5-0.5B-Instruct",
|
|
"qwen/qwen2.5-32b": "Qwen/Qwen2.5-0.5B",
|
|
"qwen/qwen2.5-32b-instruct": "Qwen/Qwen2.5-0.5B-Instruct",
|
|
"qwen/qwen3-32b": "Qwen/Qwen2.5-0.5B-Instruct",
|
|
"qwen/qwen3-14b": "Qwen/Qwen2.5-0.5B-Instruct",
|
|
# Mistral / Mixtral
|
|
"mistralai/mixtral-8x22b-instruct-v0.1": "mistralai/Mistral-7B-Instruct-v0.3",
|
|
"mistralai/mistral-large-instruct-2407": "mistralai/Mistral-7B-Instruct-v0.3",
|
|
# DeepSeek
|
|
"deepseek-ai/deepseek-v3": "deepseek-ai/DeepSeek-V3-0324",
|
|
"deepseek-ai/deepseek-r1": "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
|
|
# Gemma 2/3 (large)
|
|
"google/gemma-2-27b": "google/gemma-2-2b",
|
|
"google/gemma-2-27b-it": "google/gemma-2-2b-it",
|
|
"google/gemma-3-27b": "google/gemma-3-4b",
|
|
"google/gemma-3-27b-it": "google/gemma-3-4b-it",
|
|
}
|
|
|
|
|
|
_HF_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_\-./]*$")
|
|
|
|
|
|
def _is_valid_hf_id(name: str) -> bool:
|
|
if not name or len(name) > 200:
|
|
return False
|
|
if "\x00" in name:
|
|
return False
|
|
# Block URLs
|
|
if name.startswith(("http://", "https://", "file://")):
|
|
return False
|
|
return bool(_HF_ID_RE.match(name))
|
|
|
|
|
|
def pick_draft_model(target: str) -> Optional[str]:
|
|
"""Pick a known-good draft model for the target.
|
|
|
|
Returns None for unknown / too-small targets. Target names are normalised
|
|
to lowercase for matching. URLs and null-byte names are rejected for
|
|
defence-in-depth (the caller loads the returned value via
|
|
``AutoModelForCausalLM.from_pretrained``).
|
|
"""
|
|
if not _is_valid_hf_id(target):
|
|
return None
|
|
key = target.strip().lower()
|
|
# Strip trailing path segments (e.g. revisions) not present in map
|
|
return _DRAFT_PAIRS.get(key)
|