from __future__ import annotations import functools from typing import Any try: # must precede torch/transformers imports: it patches CUDA allocation import spaces except ImportError: # pragma: no cover - local run spaces = None import torch from transformers import AutoModelForCausalLM, AutoTokenizer BASELINE_REPO = "yava-code/Tessera-1B-Nano-Base" CONCEPT_REPO = "yava-code/Tessera-1B-Nano" EXAMPLE_PROMPTS = [ "The most common metals are", "Photosynthesis is the process used by plants", "In the nineteenth century, transport changed because", "A simple recipe for bread needs flour, water and", ] DEFAULT_MAX_NEW_TOKENS = 48 MAX_NEW_TOKENS_LIMIT = 160 DEFAULT_SEED = 7 INTRO_NOTE = ( "Two copies of the same 360M-parameter language model, trained on the same billion" " tokens in the same order. The right one also carries a small second module: it" " compresses every four tokens into one concept code, guesses the *next* concept" " code, and feeds that guess back into the decoder while it writes.\n\n" "Type a prompt and watch both write. Then look at the third column: that is the" " same concept model with its concept feedback silenced. Something changes - and" " the loss line under the box shows how much the model minds. Now the strange" " part, measured on held-out data: replace the concept feedback with a concept" " predicted from a *different* sentence, and the model does not care. It reads the" " channel constantly - silencing it hurts - but whatever it reads there is not the" " sentence. What is it reading, then? That question is the study.\n\n" "*(Same seed and prompt give bit-identical generations on every click; change the" " seed to sample a different draw.)*" ) MODELS: dict[str, tuple[AutoTokenizer, Any]] = {} def load_models(*, gpu_dtype: bool = False) -> None: """Load both checkpoints on CPU; gpu_dtype selects bf16 for the ZeroGPU move.""" if MODELS: return dtype = torch.bfloat16 if gpu_dtype else torch.float32 for key, repo, trust in ( ("base", BASELINE_REPO, False), ("concept", CONCEPT_REPO, True), ): tokenizer = AutoTokenizer.from_pretrained(repo) model = AutoModelForCausalLM.from_pretrained( repo, trust_remote_code=trust, dtype=dtype, ) model.eval() MODELS[key] = (tokenizer, model) def _device(model: Any) -> torch.device: return next(model.parameters()).device def _generate( tokenizer: AutoTokenizer, model: Any, prompt: str, max_new_tokens: int, *, seed: int, zero_feedback: bool = False, ) -> str: inputs = tokenizer(prompt, return_tensors="pt").to(_device(model)) if zero_feedback: # generate() rejects unknown kwargs, so route concept_mode through forward. model.forward = functools.partial(model.forward, concept_mode="zero") # transformers 4.x multinomial sampling draws from the global RNG (no generator # argument), so reproducibility needs a device-global seed reset per call. device = _device(model) if device.type == "cuda": torch.cuda.manual_seed(seed) else: torch.manual_seed(seed) try: with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.8, top_p=0.95, pad_token_id=tokenizer.eos_token_id, ) finally: if zero_feedback: del model.forward return tokenizer.decode( outputs[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True, ) def _prompt_loss( tokenizer: AutoTokenizer, model: Any, prompt: str, *, zero_feedback: bool = False, ) -> float: inputs = tokenizer(prompt, return_tensors="pt").to(_device(model)) with torch.no_grad(): outputs = model(**inputs, concept_mode="zero") if zero_feedback else model(**inputs) logits = outputs.logits[:, :-1, :] targets = inputs["input_ids"][:, 1:] return float( torch.nn.functional.cross_entropy( logits.reshape(-1, logits.size(-1)), targets.reshape(-1), reduction="mean", ) ) def _run(prompt: str, max_new_tokens: int, seed: int) -> tuple[str, str, str, str, str]: if set(MODELS) != {"base", "concept"}: raise RuntimeError("Models are still loading; retry in a minute.") prompt = prompt.strip() or EXAMPLE_PROMPTS[0] max_new_tokens = max(16, min(int(max_new_tokens), MAX_NEW_TOKENS_LIMIT)) seed = max(0, min(int(seed), 2**31 - 1)) base_tok, base_model = MODELS["base"] concept_tok, concept_model = MODELS["concept"] # Identical seeds across the arms make repeats bit-reproducible; per-arm offsets # only decorrelate the columns, they change nothing on a repeated click. base_text = _generate(base_tok, base_model, prompt, max_new_tokens, seed=seed) concept_text = _generate(concept_tok, concept_model, prompt, max_new_tokens, seed=seed + 1) zero_text = _generate( concept_tok, concept_model, prompt, max_new_tokens, seed=seed + 2, zero_feedback=True, ) base_loss = _prompt_loss(base_tok, base_model, prompt) concept_loss = _prompt_loss(concept_tok, concept_model, prompt) zero_loss = _prompt_loss(concept_tok, concept_model, prompt, zero_feedback=True) delta = zero_loss - concept_loss base_note = f"teacher-forced NTP loss on the prompt: {base_loss:.3f}" concept_note = ( f"teacher-forced NTP loss on the prompt: {concept_loss:.3f}; zeroing the concept" f" feedback costs {delta:+.3f} nats here - the model notices (held-out average below)" ) return base_text, base_note, concept_text, zero_text, concept_note if spaces is not None: # ZeroGPU: models live on CPU, compute moves to a shared A100 @spaces.GPU(duration=45) def _gpu_run(prompt: str, max_new_tokens: int, seed: int) -> tuple[str, str, str, str, str]: # Only the device move and the math run inside the GPU window: downloads and # CPU loading stay in the main process, or ZeroGPU kills the call on timeout. for _, model in MODELS.values(): if next(model.parameters()).device.type != "cuda": model.to("cuda") return _run(prompt, max_new_tokens, seed) def run_comparison( prompt: str, max_new_tokens: int, seed: int ) -> tuple[str, str, str, str, str]: load_models(gpu_dtype=True) return _gpu_run(prompt, max_new_tokens, seed) else: # local: plain cpu/cuda execution def run_comparison( prompt: str, max_new_tokens: int, seed: int ) -> tuple[str, str, str, str, str]: load_models() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") for _, model in MODELS.values(): model.to(device) return _run(prompt, max_new_tokens, seed)