Files
aituner/runs/fidelity-headroom/test_prefix_analysis.py

73 lines
2.3 KiB
Python

#!/usr/bin/env python3
from __future__ import annotations
import math
import sys
from pathlib import Path
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import analyze_prefixes as analysis # noqa: E402
def main() -> None:
exact, exact_source = analysis.completion_elapsed_s(
{"completed_elapsed_s": 7.25}
)
assert exact == 7.25 and exact_source == "exact_monotonic"
reconstructed, reconstructed_source = analysis.completion_elapsed_s(
{
"success": True,
"arrival_s": 2.0,
"ttft_ms": 100.0,
"tpot_ms": 10.0,
"completion_tokens": 11,
}
)
assert math.isclose(reconstructed or 0.0, 2.2)
assert reconstructed_source == "reconstructed_from_latency"
missing, missing_source = analysis.completion_elapsed_s({"success": False})
assert missing is None and missing_source == "unobserved_failure"
examples = [
analysis.PrefixExample(
cell=f"c{index}",
anchor=float(index),
cutoff_s=5.0,
tp=1,
full_elapsed_s=65.0,
feasible=label,
primary_feasible=label,
outcome=(float(index),),
instrumentation=(float(index % 2),),
completion_time_source="exact_monotonic",
)
for index, label in enumerate((0, 1, 1))
]
labels = analysis.np.asarray([0, 1, 1])
probabilities = analysis.np.asarray([0.01, 0.99, 0.60])
policy = analysis.policy_metrics(examples, labels, probabilities, 0.95)
assert policy["early_accept"] == 1
assert policy["early_reject"] == 1
assert policy["abstain_continue_full"] == 1
assert policy["false_accept"] == 0 and policy["false_reject"] == 0
assert policy["valid_zero_error_policy"]
assert policy["valid_cost_reduction_fraction"] is not None
model = analysis.fit_frozen_model(
examples,
instrumentation_aware=True,
regularization=1.0,
)
frozen_probability = analysis.predict_frozen_model(model, examples)
assert len(frozen_probability) == len(examples)
assert analysis.np.all(frozen_probability >= 0.0)
assert analysis.np.all(frozen_probability <= 1.0)
print("fidelity prefix analysis: PASS")
if __name__ == "__main__":
main()