mirror of https://github.com/razor-ai/soup.git
80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
"""v0.44.0 Part D — `soup delinearize-llama4` weight reshape stub.
|
|
|
|
Llama 4 ships with linearised expert weights that some downstream backends
|
|
expect in 3-D form. This module declares the planned reshape; live runtime
|
|
deferred to v0.44.1 (mirrors the project's stub-then-live pattern).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
from dataclasses import dataclass
|
|
from typing import List, Tuple
|
|
|
|
from soup_cli.utils.paths import is_under_cwd
|
|
|
|
# Restrict to canonical Llama 4 model id shape; reject crafted names.
|
|
_LLAMA4_RE = re.compile(r"(?i)(?:^|[^a-z0-9])llama-?4(?:[^a-z0-9]|$)")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class DelinearizePlan:
|
|
"""Planned weights to reshape, source path, target path.
|
|
|
|
`weight_files` is a `tuple` for genuine immutability (matches the
|
|
project frozen-collection policy).
|
|
"""
|
|
|
|
source_dir: str
|
|
target_dir: str
|
|
weight_files: Tuple[str, ...]
|
|
|
|
|
|
def is_llama4_model(name: str) -> bool:
|
|
"""Return True iff `name` looks like a Llama 4 family model."""
|
|
if not isinstance(name, str) or not name or "\x00" in name:
|
|
return False
|
|
return bool(_LLAMA4_RE.search(name))
|
|
|
|
|
|
def discover_weight_files(source_dir: str) -> List[str]:
|
|
"""List `.safetensors` weight files in `source_dir`."""
|
|
if not isinstance(source_dir, str):
|
|
raise TypeError("source_dir must be str")
|
|
if not is_under_cwd(source_dir):
|
|
raise ValueError(
|
|
f"source_dir is outside cwd: {os.path.basename(source_dir)}"
|
|
)
|
|
real = os.path.realpath(source_dir)
|
|
if not os.path.isdir(real):
|
|
raise FileNotFoundError(
|
|
f"source_dir not found: {os.path.basename(real)}"
|
|
)
|
|
files = sorted(
|
|
entry for entry in os.listdir(real) if entry.endswith(".safetensors")
|
|
)
|
|
if not files:
|
|
raise FileNotFoundError(
|
|
"no .safetensors files found in source_dir"
|
|
)
|
|
return files
|
|
|
|
|
|
def plan_delinearize(source_dir: str, target_dir: str) -> DelinearizePlan:
|
|
"""Build a `DelinearizePlan`. Raises on bad inputs."""
|
|
if not isinstance(target_dir, str) or not target_dir:
|
|
raise ValueError("target_dir must be non-empty str")
|
|
if "\x00" in target_dir:
|
|
raise ValueError("target_dir contains NUL byte")
|
|
if not is_under_cwd(target_dir):
|
|
raise ValueError(
|
|
f"target_dir is outside cwd: {os.path.basename(target_dir)}"
|
|
)
|
|
files = discover_weight_files(source_dir)
|
|
return DelinearizePlan(
|
|
source_dir=os.path.realpath(source_dir),
|
|
target_dir=os.path.realpath(target_dir),
|
|
weight_files=tuple(files),
|
|
)
|