diff --git a/recipe/jemm/export.py b/recipe/jemm/export.py index 21f05b20..f0f599ce 100644 --- a/recipe/jemm/export.py +++ b/recipe/jemm/export.py @@ -31,8 +31,6 @@ LABELS = "ABCDEFGHIJKLMNOPQRSTUVWXYZ012345" SYSTEM = ("Choose the best available candidate for the question using only the supplied state. " "Return exactly one candidate label.") -CALIBRATION = {"temperature": 1.3480874159655591, "mm_temperature": 1.3954832341582943, - "threshold": 0.9872681877423998} INVENTORY_PATH = Path(__file__).with_name("pinned_inventory.json") @@ -189,9 +187,14 @@ def export_checkpoint(base, adapter, out, provenance, *, tokenizer=None, invento pairs, scale = adapter_pairs(inventory["base"]["tensors"], inventory["adapter"]["tensors"], adapter_config) if inventory.get("provenance") == PINS and (adapter_config["r"] != 16 or len(pairs) != 496 or scale != 2.): raise ValueError("pinned JEMM requires rank16, 496 FP32 LoRA pairs and scale2") - calibration = json.loads((adapter / "decision_config.json").read_text()) - if any(calibration.get(key) != value for key, value in CALIBRATION.items()): - raise ValueError("decision calibration differs from pinned JEMM") + decision = json.loads((adapter / "decision_config.json").read_text()) + calibration = {} + for key in ("temperature", "mm_temperature", "threshold"): + value = decision.get(key) if isinstance(decision, dict) else None + if type(value) not in (int, float) or not math.isfinite(value) or ( + value <= 0 if key != "threshold" else not 0 <= value <= 1): + raise ValueError(f"invalid decision calibration: {key}") + calibration[key] = value if tokenizer is None: from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(str(base), local_files_only=True) @@ -310,7 +313,7 @@ def export_checkpoint(base, adapter, out, provenance, *, tokenizer=None, invento raise ValueError("original FP32 adapter checksum mismatch") if "jemm_lm_head.safetensors" not in exports or "vision.safetensors" not in exports: raise ValueError("export lacks selected head or vision tensors") - manifest = {"format": "jemm-native/2", "model_id": "JEMM", **PINS, **CALIBRATION, + manifest = {"format": "jemm-native/2", "model_id": "JEMM", **PINS, **calibration, "max_tokens": 8192, "max_mm_tokens": 3072, "chat_prefix": chat_prefix, "chat_suffix": chat_suffix, "enable_thinking": False, "preserve_thinking": False, "labels": list(LABELS), "label_token_ids": label_ids, "label_head_file": "jemm_lm_head.safetensors", "label_head_tensor": "weight", "label_head_dtype": "BF16", diff --git a/src/models/jemm/README.md b/src/models/jemm/README.md index 5238ebc1..ccd77bfe 100644 --- a/src/models/jemm/README.md +++ b/src/models/jemm/README.md @@ -21,10 +21,11 @@ rows from the untied LM head are extracted without merging adapter weights. `jemm_export.json` records every exported file's SHA-256 and the adapter contract. Startup verifies original shard and adapter pins before CUDA initialization; legacy premerged `jemm-native/1` exports are explicitly rejected. Keep files immutable -while serving. This head requires the pinned temperatures 1.3480874159655591 -(text), 1.3954832341582943 (images) and threshold 0.9872681877423998. Follow-up -#122 loads validated calibration from the checksum-covered decision config; -its calibration policy is a separate change. The threshold is provenance and +while serving. Calibration loads from checksum-covered `decision_config.json`; +temperatures must be finite and positive, threshold finite and within [0,1], +and all values must agree with the export manifest. The pinned artifact values +are 1.3480874159655591 (text), 1.3954832341582943 (images) and +0.9872681877423998 (threshold). The threshold is provenance and does not gate SystemOne responses. ## Request and response contract diff --git a/src/models/jemm/native/src/artifacts.rs b/src/models/jemm/native/src/artifacts.rs index f687dcb0..f3fd7847 100644 --- a/src/models/jemm/native/src/artifacts.rs +++ b/src/models/jemm/native/src/artifacts.rs @@ -104,6 +104,7 @@ pub fn verify_export_files(dir: &Path, manifest: &Value) -> Result<()> { "config.json", "tokenizer.json", "preprocessor_config.json", + "decision_config.json", "model.safetensors.index.json", "vision.safetensors", "jemm_lm_head.safetensors", diff --git a/src/models/jemm/native/src/contract.rs b/src/models/jemm/native/src/contract.rs index a63147d0..ec26403d 100644 --- a/src/models/jemm/native/src/contract.rs +++ b/src/models/jemm/native/src/contract.rs @@ -11,9 +11,6 @@ pub const LABELS: &str = "ABCDEFGHIJKLMNOPQRSTUVWXYZ012345"; pub const SYSTEM: &str = "Choose the best available candidate for the question using only the supplied state. Return exactly one candidate label."; pub const BASE_REVISION: &str = "1d4bf0f2ff6012fd82039f2fa52739d0dd7c60c0"; pub const CHECKPOINT_REVISION: &str = "76e3c209e8441fa658221c7ba2725bad2f811176"; -pub const TEMPERATURE: f64 = 1.3480874159655591; -pub const MM_TEMPERATURE: f64 = 1.3954832341582943; -pub const THRESHOLD: f64 = 0.9872681877423998; pub const SOURCE_REVISION: &str = "6822fe0fd53c5e6670af6ba99fb2c857a661e532"; #[derive(Debug)] pub struct Question { diff --git a/src/models/jemm/native/src/processing.rs b/src/models/jemm/native/src/processing.rs index 3538da9f..759b811d 100644 --- a/src/models/jemm/native/src/processing.rs +++ b/src/models/jemm/native/src/processing.rs @@ -200,16 +200,33 @@ impl Processor { prefix.contains(contract::SYSTEM) && suffix.contains("\n\n"), "expected JEMM system and disabled-thinking template" ); - let temperature = manifest["temperature"].as_f64().context("temperature")?; - let mm_temperature = manifest["mm_temperature"] - .as_f64() - .context("mm_temperature")?; - ensure!( - temperature == contract::TEMPERATURE - && mm_temperature == contract::MM_TEMPERATURE - && manifest["threshold"].as_f64() == Some(contract::THRESHOLD), - "expected pinned JEMM calibration" - ); + let calibration: Value = serde_json::from_slice( + &std::fs::read(dir.join("decision_config.json")).context("decision_config.json")?, + )?; + let mut values = [0.0; 3]; + for (index, key) in ["temperature", "mm_temperature", "threshold"] + .iter() + .enumerate() + { + let value = calibration[key] + .as_f64() + .with_context(|| format!("calibration {key}"))?; + ensure!( + value.is_finite() + && if index < 2 { + value > 0.0 + } else { + (0.0..=1.0).contains(&value) + }, + "invalid calibration {key}" + ); + ensure!( + manifest[key].as_f64() == Some(value), + "manifest disagrees with decision_config.json: {key}" + ); + values[index] = value; + } + let [temperature, mm_temperature, _threshold] = values; let config: Value = serde_json::from_slice(&std::fs::read(dir.join("config.json"))?)?; let image_token = config["image_token_id"] .as_u64() diff --git a/tests/jemm/test_export.py b/tests/jemm/test_export.py index 75a652dd..c4997920 100644 --- a/tests/jemm/test_export.py +++ b/tests/jemm/test_export.py @@ -62,7 +62,9 @@ def setUp(self): self.write_adapter() self.config = {"peft_type": "LORA", "base_model_name_or_path": EXPORT.BASE_MODEL_ID, "r": 16, "lora_alpha": 32} (self.adapter / "adapter_config.json").write_text(json.dumps(self.config)) - (self.adapter / "decision_config.json").write_text(json.dumps(EXPORT.CALIBRATION)) + (self.adapter / "decision_config.json").write_text(json.dumps({"temperature": 1.3480874159655591, + "mm_temperature": 1.3954832341582943, + "threshold": 0.9872681877423998})) self.refresh_inventory() def tearDown(self): @@ -103,6 +105,31 @@ def test_preserves_original_base_and_fp32_adapter_without_merging(self): with safe_open(str(self.out / "vision.safetensors"), framework="pt", device="cpu") as reader: self.assertTrue(torch.equal(reader.get_tensor("model.visual.patch_embed.weight"), torch.ones(2, 2, dtype=torch.bfloat16))) + def test_v2_export_uses_checksum_covered_file_calibration(self): + calibration = {"temperature": 2., "mm_temperature": 3., "threshold": 0.5} + source = self.adapter / "decision_config.json" + source.write_text(json.dumps(calibration)) + self.refresh_inventory() + manifest = self.run_export() + self.assertEqual(manifest["format"], "jemm-native/2") + self.assertEqual(manifest["lora_mode"], "unmerged_fp32") + self.assertEqual({key: manifest[key] for key in calibration}, calibration) + self.assertEqual((self.out / source.name).read_bytes(), source.read_bytes()) + self.assertEqual(manifest["export_sha256"][source.name], EXPORT.sha256(source)) + + def test_invalid_file_calibration_is_refused_before_output_creation(self): + for key, value in [("temperature", 0), ("temperature", True), + ("mm_temperature", -1), ("threshold", 1.1), + ("threshold", float("nan")), ("temperature", None)]: + with self.subTest(key=key, value=value): + calibration = {"temperature": 2., "mm_temperature": 3., "threshold": 0.5} + calibration[key] = value + (self.adapter / "decision_config.json").write_text(json.dumps(calibration)) + self.refresh_inventory() + with self.assertRaisesRegex(ValueError, "invalid decision calibration"): + self.run_export() + self.assertFalse(self.out.exists()) + def test_cross_device_copy_fallback_preserves_identical_source_bytes(self): with patch.object(EXPORT.os, "link", side_effect=OSError(errno.EXDEV, "cross-device")): manifest = self.run_export() diff --git a/tests/jemm/tokenization.rs b/tests/jemm/tokenization.rs index f9bc30f6..8aaca82e 100644 --- a/tests/jemm/tokenization.rs +++ b/tests/jemm/tokenization.rs @@ -43,6 +43,13 @@ fn fixture() -> (Fixture, Value) { ) .unwrap(); let manifest = json!({"format":"jemm-native/2","model_id":"JEMM","base_model_id":"Qwen/Qwen3.8-27B","base_revision":contract::BASE_REVISION,"checkpoint_revision":contract::CHECKPOINT_REVISION,"source_revision":contract::SOURCE_REVISION,"max_tokens":8192,"max_mm_tokens":3072,"label_token_ids":(1..=32).collect::>(),"chat_prefix":format!("<|im_start|>system\n{}<|im_end|>\n<|im_start|>user\n",contract::SYSTEM),"chat_suffix":"<|im_end|>\n<|im_start|>assistant\n\n\n\n\n","temperature":1.3480874159655591,"mm_temperature":1.3954832341582943,"threshold":0.9872681877423998}); + std::fs::write( + dir.join("decision_config.json"), + serde_json::to_vec(&json!({"temperature":manifest["temperature"], + "mm_temperature":manifest["mm_temperature"], "threshold":manifest["threshold"]})) + .unwrap(), + ) + .unwrap(); (Fixture(dir), manifest) } #[test] @@ -86,7 +93,7 @@ fn exported_tokenizer_truncation_and_padding_cannot_change_request_tokens() { assert!(prepared.inputs[0].ids.len() > 3); } #[test] -fn rejects_custom_calibration_for_pinned_adapter() { +fn rejects_manifest_calibration_disagreement() { let (f, mut m) = fixture(); m["mm_temperature"] = json!(1.); assert!( @@ -94,7 +101,52 @@ fn rejects_custom_calibration_for_pinned_adapter() { .err() .unwrap() .to_string() - .contains("pinned JEMM calibration") + .contains("manifest disagrees with decision_config.json") + ); +} + +#[test] +fn file_driven_calibration_must_be_finite_in_range_and_agree_with_manifest() { + for (field, invalid) in [ + ("temperature", json!(0.0)), + ("mm_temperature", json!(-1.0)), + ("threshold", json!(1.1)), + ] { + let (f, mut m) = fixture(); + m[field] = invalid; + std::fs::write( + f.0.join("decision_config.json"), + serde_json::to_vec(&json!({"temperature":m["temperature"], + "mm_temperature":m["mm_temperature"], "threshold":m["threshold"]})) + .unwrap(), + ) + .unwrap(); + assert!( + Processor::load(&f.0, &m) + .err() + .unwrap() + .to_string() + .contains("invalid calibration") + ); + } + let (f, mut m) = fixture(); + m["temperature"] = json!(2.0); + m["mm_temperature"] = json!(3.0); + m["threshold"] = json!(0.5); + std::fs::write( + f.0.join("decision_config.json"), + serde_json::to_vec(&json!({"temperature":2.0,"mm_temperature":3.0,"threshold":0.5})) + .unwrap(), + ) + .unwrap(); + let p = Processor::load(&f.0, &m).unwrap(); + let prepared = p + .prepare(br#"{"questions":{"q":{"type":"noul"}}}"#) + .unwrap(); + let response = prepared.context.finish(vec![vec![2.0, 0.0]]).unwrap(); + assert!( + (response["answers"]["q"]["noul"].as_f64().unwrap() - 1.0 / (1.0 + (-1.0f64).exp())).abs() + < 1e-12 ); }