Checkpoints from the
Cocoracle experiment -- interpreting what a model "thinks" during latent reasoning.
Combines
Coconut (Chain of Continuous Thought) with
Activation Oracles to train models that answer natural-language questions about their own latent chain-of-thought hidden states.
1import torch
2from transformers import GPT2LMHeadModel, GPT2Tokenizer
3
4tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
5tokenizer.pad_token = tokenizer.eos_token
6tokenizer.add_special_tokens({
7 "additional_special_tokens": ["<bot>", "<sep>", "<eot>", "<act>"]
8})
9
10model = GPT2LMHeadModel.from_pretrained("gpt2-large")
11model.resize_token_embeddings(len(tokenizer))
12state = torch.load("stage3_alllatent.pt", map_location="cpu")
13model.load_state_dict(state)
See the
GitHub repo for full code and an interactive demo (
scripts/interactive.py).