Views
No views yet
<|reserved_special_token_12|> tokens. The standard transformers library won't handle the continuous token embeddings correctly.1import torch
2import numpy as np
3from transformers import AutoTokenizer
4
5# Load the continuous model class
6from model.continuous_llama import ContinuousLlama
7
8# Load the model and tokenizer
9model_name = "Transluce/features_explain_llama3.1_8b_llama3.1_8b_instruct"
10tokenizer = AutoTokenizer.from_pretrained(model_name)
11model = ContinuousLlama.from_pretrained(
12 model_name,
13 torch_dtype=torch.bfloat16,
14 special_tokens_ids={
15 "begin_continuous": tokenizer.convert_tokens_to_ids("<|reserved_special_token_10|>"),
16 "end_continuous": tokenizer.convert_tokens_to_ids("<|reserved_special_token_11|>"),
17 "continuous_rep": tokenizer.convert_tokens_to_ids("<|reserved_special_token_12|>")
18 }
19)
20
21# Example: explaining a continuous feature from layer 15
22layer = 15
23feature_vector = torch.randn(4096) # Feature from Llama-3.1-8B's residual stream
24
25# Format the prompt with continuous tokens
26prompt = [{
27 "role": "user",
28 "content": f"At layer {layer}, <|reserved_special_token_10|><|reserved_special_token_12|><|reserved_special_token_11|> encodes "
29}]
30chat_prompt = tokenizer.apply_chat_template(prompt, tokenize=False)
31
32# Tokenize the prompt
33inputs = tokenizer(prompt, return_tensors="pt")
34
35# Create continuous token inputs for the feature vector
36continuous_tokens = {
37 "inputs_continuous_tokens": feature_vector.unsqueeze(0), # Add batch dimension
38 "labels_continuous_tokens": None # Not needed for generation
39}
40
41# Generate explanation
42with torch.no_grad():
43 outputs = model.generate(
44 input_ids=inputs.input_ids,
45 attention_mask=inputs.attention_mask,
46 max_new_tokens=128,
47 do_sample=False,
48 pad_token_id=tokenizer.eos_token_id,
49 **continuous_tokens
50 )
51
52# Decode the explanation
53explanation = tokenizer.decode(outputs[0], skip_special_tokens=True)
54print(explanation)@misc{li2025traininglanguagemodelsexplain,
title={Training Language Models to Explain Their Own Computations},
author={Belinda Z. Li and Zifan Carl Guo and Vincent Huang and Jacob Steinhardt and Jacob Andreas},
year={2025},
eprint={2511.08579},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2511.08579},
}