ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter8/prompt-distillation/test_experiment_8_8.py

@@ -0,0 +1,65 @@
+import importlib.util
+import sys
+from pathlib import Path
+
+
+MODULE_PATH = Path(__file__).with_name("run_experiment_8_8.py")
+SPEC = importlib.util.spec_from_file_location("experiment_8_8", MODULE_PATH)
+exp = importlib.util.module_from_spec(SPEC)
+assert SPEC and SPEC.loader
+sys.modules[SPEC.name] = exp
+SPEC.loader.exec_module(exp)
+
+
+def test_parse_label_is_strict():
+    assert exp.parse_label("Final Answer: fr") == "fr"
+    assert exp.parse_label("zh") == "zh"
+    assert exp.parse_label("Final Answer: jp") is None
+    assert exp.parse_label("The answer is French") is None
+
+
+def test_cost_uses_uncached_cached_and_output_rates():
+    usage = {
+        "prompt_tokens": 100,
+        "completion_tokens": 10,
+        "prompt_tokens_details": {"cached_tokens": 40},
+    }
+    cost = exp.usage_cost(usage)
+    expected_cny = (60 * 20 + 40 * 2 + 10 * 100) / 1_000_000
+    assert cost["cny"] == expected_cny
+    assert cost["usd"] == expected_cny * exp.PRICING["usd_per_currency_unit"]
+
+
+def test_stable_sample_maps_out_of_scope_languages_to_other():
+    rows = [
+        {"labels": "fr", "text": "bonjour"},
+        {"labels": "ja", "text": "こんにちは"},
+    ]
+    sampled = exp._stable_sample(rows, n=1, seed=1, split="train")
+    by_source = {row["source_label"]: row for row in sampled}
+    assert by_source["fr"]["gold_label"] == "fr"
+    assert by_source["ja"]["gold_label"] == "ot"
+
+
+def test_sft_rows_expose_only_raw_text_to_student():
+    row = {
+        "text": "bonjour",
+        "prediction": "fr",
+        "id": "x",
+        "gold_label": "fr",
+        "response_id": "receipt",
+    }
+    sample = {
+        "id": row["id"],
+        "messages": [
+            {"role": "user", "content": row["text"]},
+            {"role": "assistant", "content": row["prediction"]},
+        ],
+    }
+    assert exp.LANGUAGE_CLASSIFICATION_PROMPT not in sample["messages"][0]["content"]
+    assert sample["messages"][0]["content"] == "bonjour"
+
+
+def test_chat_template_ids_accepts_transformers_4_and_5_shapes():
+    assert exp._chat_template_ids([1, 2, 3]) == [1, 2, 3]
+    assert exp._chat_template_ids({"input_ids": [1, 2, 3], "attention_mask": [1, 1, 1]}) == [1, 2, 3]