"""Exercise real attention extraction with a tiny synthetic, untrained test model.""" import json import tempfile import unittest from pathlib import Path from unittest.mock import patch import torch from fastapi.testclient import TestClient from tokenizers import Regex, Tokenizer from tokenizers.models import WordLevel from tokenizers.pre_tokenizers import Split from transformers import ( GraniteMoeHybridConfig, GraniteMoeHybridForCausalLM, LlamaConfig, LlamaForCausalLM, PreTrainedTokenizerFast, ) import app as demo from runtime import EXAMPLES, MODEL_ID, AttentionRequest, InferenceEngine def tiny_model(path: Path) -> None: """Save a tiny random model for integration tests, never for deployment.""" alphabet = sorted(set("".join(ex["llmon_prompt_dirty"] for ex in EXAMPLES.values()))) vocab = {"[UNK]": 0, "[EOS]": 1, **{c: i + 2 for i, c in enumerate(alphabet)}} tokenizer = Tokenizer(WordLevel(vocab, unk_token="[UNK]")) tokenizer.pre_tokenizer = Split(Regex(""), behavior="isolated") fast = PreTrainedTokenizerFast(tokenizer_object=tokenizer, unk_token="[UNK]", eos_token="[EOS]") fast.save_pretrained(path) torch.manual_seed(7) config = LlamaConfig( vocab_size=len(vocab), hidden_size=16, intermediate_size=32, num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=1, max_position_embeddings=2048, eos_token_id=1, ) LlamaForCausalLM(config).save_pretrained(path) class LiveRuntimeTests(unittest.TestCase): """Test the hosted adapter without claiming access to the deployed weights.""" @classmethod def setUpClass(cls) -> None: """Load the small real Transformers model on CPU for portable checks.""" cls.temp = tempfile.TemporaryDirectory() cls.path = Path(cls.temp.name) tiny_model(cls.path) cls.engine = InferenceEngine() with patch.dict("os.environ", {"LLMON_MODEL_PATH": str(cls.path), "LLMON_DEVICE": "cpu", "SPACE_ID": ""}): cls.engine.prepare() cls.client = TestClient(demo.app) @classmethod def tearDownClass(cls) -> None: """Remove the test-only checkpoint.""" cls.temp.cleanup() def test_current_examples_and_actual_attention(self) -> None: """Check all three new prompts, target spans, normalized metrics, and layers.""" for example_id in EXAMPLES: with self.subTest(example_id=example_id): result = self.engine.infer("prefill", {"example_id": example_id}) self.assertEqual(result["model"]["model_id"], MODEL_ID) self.assertEqual(len(result["tokens"]), len(result["span_labels"])) self.assertTrue(any(s["is_target"] and s["type"] == "instr" for s in result["spans"])) self.assertEqual(set(result["attention_by_layer"]), {"0", "1"}) for layer in result["attention_by_layer"].values(): for metric in layer["metrics"].values(): self.assertGreaterEqual(metric, 0) self.assertLessEqual(metric, 1) json.dumps(result, allow_nan=False) def test_generation_and_benchmark(self) -> None: """Run generation on a requested layer and compute current benchmark rows.""" result = self.engine.infer("generate", {"example_id": 0, "layer": 0, "max_tokens": 2}) self.assertEqual(result["layer_used"], 0) self.assertTrue(1 <= len(result["generated_tokens"]) <= 2) self.assertEqual(len(result["generated_tokens"]), len(result["attention_per_step"])) self.assertEqual(len(result["generated_token_ids"]), len(result["generated_tokens"])) self.assertEqual( result["generated_text"], self.engine.tokenizer.decode( result["generated_token_ids"], skip_special_tokens=True, clean_up_tokenization_spaces=False ), ) self.assertEqual( result["generation_stop_reason"], "eos" if result["generated_token_ids"][-1] == self.engine.tokenizer.eos_token_id else "token_limit", ) rows = self.engine.benchmark()["results"] self.assertEqual({row["example_id"] for row in rows}, set(EXAMPLES)) self.assertTrue(all(row["model_id"] == MODEL_ID for row in rows)) json.dumps(rows, allow_nan=False) def test_granite_hybrid_attention_layer_mapping(self) -> None: """Preserve real layer indices across Granite attention and Mamba blocks.""" with tempfile.TemporaryDirectory() as directory: path = Path(directory) self.engine.tokenizer.save_pretrained(path) config = GraniteMoeHybridConfig( vocab_size=len(self.engine.tokenizer), hidden_size=16, intermediate_size=32, shared_intermediate_size=32, num_hidden_layers=3, num_attention_heads=2, num_key_value_heads=1, layer_types=["attention", "mamba", "attention"], num_local_experts=0, num_experts_per_tok=0, mamba_n_heads=2, mamba_d_head=16, mamba_d_state=4, mamba_chunk_size=32, max_position_embeddings=2048, position_embedding_type="rope", bos_token_id=0, eos_token_id=1, ) torch.manual_seed(7) GraniteMoeHybridForCausalLM(config).save_pretrained(path) engine = InferenceEngine() with patch.dict("os.environ", {"LLMON_MODEL_PATH": str(path), "LLMON_DEVICE": "cpu", "SPACE_ID": ""}): engine.prepare() result = engine.infer("prefill", {"example_id": 0}) layers = result["attention_by_layer"] self.assertEqual(set(layers), {"0", "1", "2"}) self.assertEqual(layers["1"]["layer_type"], "mamba") self.assertIsNone(layers["1"]["metrics"]) self.assertIsNotNone(layers["2"]["metrics"]) generated = engine.infer("generate", {"example_id": 0, "layer": 2, "max_tokens": 2}) self.assertEqual(generated["layer_used"], 2) self.assertGreaterEqual(len(generated["generated_tokens"]), 1) with self.assertRaises(ValueError): engine.infer("generate", {"example_id": 0, "layer": 1}) json.dumps(result, allow_nan=False) json.dumps(generated, allow_nan=False) def test_limits_and_invalid_requests(self) -> None: """Reject paths, unsupported layers, ambiguous prompts, and oversized input.""" invalid = [ {}, {"example_id": 99}, {"example_id": 0, "model_id": "../../etc"}, {"example_id": 0, "max_tokens": 1000}, {"example_id": 0, "prompt_llmon": "x"}, {"example_id": 0, "dataset_tag": "../../"}, {"prompt_llmon": "x" * 20000}, {"example_id": 0, "model_path": "/tmp/model"}, ] for payload in invalid: with self.subTest(payload=list(payload)), self.assertRaises(ValueError): AttentionRequest.model_validate(payload) with self.assertRaises(ValueError): self.engine.infer("generate", {"example_id": 0, "layer": 255}) with patch.object(self.engine, "max_input_tokens", 8), self.assertRaises(ValueError): self.engine.infer("prefill", {"example_id": 0}) with self.assertRaises(ValueError): self.engine.infer("prefill", {"prompt_llmon": "plain text without LLMON tags"}) def test_missing_model_never_substitutes_recordings(self) -> None: """Keep the app available while explicitly reporting missing live inference.""" with patch.object(demo, "engine", InferenceEngine()): health = self.client.get("/api/health").json() self.assertFalse(health["model_ready"]) self.assertEqual(health["mode"], "live") self.assertEqual(demo.prefill({"example_id": 0})["error"]["status"], 503) def test_no_admin_or_model_files_are_served(self) -> None: """Publish the landing, guide, dashboards and intended APIs.""" for path in [ "/", "/how-it-works/", "/experiment_1/", "/experiment_2/", "/generation-flow-2/", "/experiment_3/", ]: self.assertEqual(self.client.get(path).status_code, 200) for path in ["/admin/", "/models/", "/vendor/", "/new_dlb_examples.jsonl", "/.env"]: self.assertEqual(self.client.get(path).status_code, 404) for path in ["/api/models/delete", "/api/models/import_hf", "/api/load_model", "/api/uploads/jsonl"]: self.assertEqual(self.client.post(path, json={}).status_code, 404) self.assertEqual(len(self.client.get("/api/data/examples").json()), 3) def test_old_comparison_links_open_the_selected_flow(self) -> None: """Retire the comparison bundle while retaining existing shared links.""" for path in ["/generation-flow-2", "/generation-flow-2/", "/generation-flow-2/index.html"]: response = self.client.get(path, follow_redirects=False) self.assertEqual(response.status_code, 307) self.assertEqual(response.headers["location"], "/experiment_2/index.html") page = self.client.get("/experiment_2/index.html").text self.assertIn("Generation Flow — LLMON", page) self.assertNotIn("Generation Flow 2", page) # The Flow page chrome moved from the page's inline stylesheet into the # design override layer; the page must load it, and it must be served. self.assertIn('href="/design-assets/css/design-overrides.css"', page) self.assertIn(".flow-summary-grid", self.client.get("/design-assets/css/design-overrides.css").text) self.assertEqual(self.client.get("/generation-flow-2/app.js").status_code, 404) if __name__ == "__main__": unittest.main()