Spaces:
Running on Zero
Running on Zero
File size: 10,158 Bytes
5c49122 9aa89b1 5c49122 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 | """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()
|