Spaces:
Running on Zero
Running on Zero
| """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.""" | |
| 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) | |
| 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() | |