Data-Constrained Scaling Law

from __future__ import annotations
import math
from typing import List, Dict

FEATURES = ['log_params', 'log_tokens', 'log_unique_tokens', 'log_params:log_tokens', 'tokens_inv_sqrt']
GROUP_PARAMS = {
  "all_data": {
    "intercept": 21.017514457355812,
    "coefs": [
      -0.8745591289420206,
      -0.5044278099541184,
      -0.11831988273483411,
      0.029264245542033336,
      38424.77315308764
    ]
  }
}
GLOBAL_PARAMS = {
  "intercept": 21.017514457355812,
  "coefs": [
    -0.8745591289420206,
    -0.5044278099541184,
    -0.11831988273483411,
    0.029264245542033336,
    38424.77315308764
  ]
}

def _feature_vector(params: float, tokens: float, unique_tokens: float):
    lp = math.log(params)
    lt = math.log(tokens)
    lu = math.log(unique_tokens)
    vec = []
    for feat in FEATURES:
        if feat == "log_params":
            vec.append(lp)
        elif feat == "log_tokens":
            vec.append(lt)
        elif feat == "log_unique_tokens":
            vec.append(lu)
        elif feat == "log_params:log_tokens":
            vec.append(lp*lt)
        elif feat == "tokens_inv_sqrt":
            vec.append(1.0/math.sqrt(tokens))
        else:
            vec.append(0.0)
    return vec

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values.
        group: The name of the experimental group for which to make predictions.
                The functional form of the law must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """
    coeffs = GROUP_PARAMS.get(group, GLOBAL_PARAMS)
    intercept = coeffs["intercept"]
    beta = coeffs["coefs"]
    out: list[dict[str, float]] = []
    for row in input_data:
        p = float(row["params"])
        t = float(row["tokens"])
        u = float(row["unique_tokens"])
        x = _feature_vector(p, t, u)
        y = intercept + sum(b*xi for b, xi in zip(beta, x))
        out.append({"loss": float(y)})
    return out

from math import log, exp, isfinite

# Fitted coefficients for log-linear power-law model:
# log(loss) = b0 + bP*log(params) + bT*log(tokens) + bU*log(unique_tokens)
_COEFS = {
  "all_data": [
    4.489044805418067,
    -0.06713156032896134,
    -0.057418372927797716,
    -0.02821632111651312
  ]
}
_GLOBAL = [4.489044805418067, -0.06713156032896134, -0.057418372927797716, -0.02821632111651312]


def _predict_one(x: dict[str, float], b: list[float]) -> float:
    P = float(x.get('params', 0.0) or 0.0)
    T = float(x.get('tokens', 0.0) or 0.0)
    U = float(x.get('unique_tokens', 0.0) or 0.0)
    if not (P > 0 and T > 0 and U > 0 and isfinite(P) and isfinite(T) and isfinite(U)):
        return float('nan')
    b0, bP, bT, bU = b
    return exp(b0 + bP*log(P) + bT*log(T) + bU*log(U))


def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values.
        group: The name of the experimental group for which to make predictions.
                The functional form of the law must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """
    b = _COEFS.get(group)
    if not (isinstance(b, list) and len(b) == 4):
        b = _GLOBAL
    return [{"loss": _predict_one(x, b)} for x in input_data]

# Auto-generated scaling law implementation
from __future__ import annotations
import math

# Per-group coefficients for: loss = a + b*ln(params) + c*ln(tokens) + d*ln(unique_tokens)
COEFS = {
  "all_data": {
    "a": 17.069182044828956,
    "b": -0.3028327010879326,
    "c": -0.2721077163179839,
    "d": -0.05760888894597566
  }
}

DEFAULT_GROUP = list(COEFS.keys())[0] if COEFS else 'default'

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values.
        group: The name of the experimental group for which to make predictions.
                The functional form of the law must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """
    coeffs = COEFS.get(group)
    if coeffs is None:
        coeffs = COEFS.get(DEFAULT_GROUP, {'a':0.0,'b':0.0,'c':0.0,'d':0.0})
    a = float(coeffs['a']); b = float(coeffs['b']); c = float(coeffs['c']); d = float(coeffs['d'])
    out: list[dict[str, float]] = []
    for row in input_data:
        p = float(row.get('params', 0.0))
        t = float(row.get('tokens', 0.0))
        u = float(row.get('unique_tokens', 0.0))
        if p <= 0 or t <= 0 or u <= 0:
            y = float('nan')
        else:
            y = a + b * math.log(p) + c * math.log(t) + d * math.log(u)
        out.append({'loss': float(y)})
    return out

from __future__ import annotations
from math import pow
from typing import Dict, List

# Shared exponents across groups
_ALPHA = 2.88154375571247
_BETA = 0.4714873693356799
_GAMMA = 2.468731307378303

# Per-group coefficients
_COEFFS: Dict[str, Dict[str, float]] = {'all_data': {'L_inf': 2.9842494330943747, 'A': 0.0, 'B': 27959.06945832133, 'C': 8.297512369762968e-12}}

# Fallback group (if an unknown group is requested)
_FALLBACK_GROUP = next(iter(_COEFFS.keys())) if _COEFFS else 'default'

def _get_group(g: str) -> str:
    return g if g in _COEFFS else _FALLBACK_GROUP

def law(input_data: list[dict[str, float]], group: str) -> list[dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values.
        group: The name of the experimental group for which to make predictions.
                The functional form of the law must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """
    g = _get_group(group)
    c = _COEFFS[g]
    a, b, cexp = _ALPHA, _BETA, _GAMMA
    out: List[Dict[str, float]] = []
    for row in input_data:
        P = float(row.get('params', 0.0))
        T = float(row.get('tokens', 0.0))
        U = float(row.get('unique_tokens', 0.0))
        # Guard against non-positive inputs
        if P <= 0 or T <= 0 or U <= 0:
            P = max(P, 1e-12)
            T = max(T, 1e-12)
            U = max(U, 1e-12)
        pred = (
            c['L_inf']
            + c['A'] * pow(P, -a)
            + c['B'] * pow(T, -b)
            + c['C'] * pow(U, -cexp)
        )
        out.append({'loss': float(pred)})
    return out