In the AxBench paper, we finetuned a subspace generator. The subspace generator is a hyper-network that will generate a subspace for you given a concept description in natural language. High-quality subspace generator can bypass all dictionary training!
1import torch
2import torch.nn.functional as F
3from transformers import AutoModelForCausalLM, AutoTokenizer
4
5class RegressionWrapper(torch.nn.Module):
6 def __init__(self, base_model, hidden_size, output_dim):
7 super().__init__()
8 self.base_model = base_model
9 self.regression_head = torch.nn.Linear(hidden_size, output_dim)
10
11 def forward(self, input_ids, attention_mask):
12 outputs = self.base_model.model(
13 input_ids=input_ids,
14 attention_mask=attention_mask,
15 output_hidden_states=True,
16 return_dict=True
17 )
18 last_hiddens = outputs.hidden_states[-1]
19 last_token_representations = last_hiddens[:, -1]
20 preds = self.regression_head(last_token_representations)
21 preds = F.normalize(preds, p=2, dim=-1)
22 return preds
23
24base_model = AutoModelForCausalLM.from_pretrained(
25 f"google/gemma-2-2b", torch_dtype=torch.bfloat16)
26base_tokenizer = AutoTokenizer.from_pretrained(
27 f"google/gemma-2-2b", model_max_length=512)
28
29subspace_gen = RegressionWrapper(
30 base_model, hidden_size, output_dim).bfloat16().to("cuda")
31subspace_gen.load_state_dict(torch.load('model.pth'))
32
33your_new_concept = "terms related to Stanford University"
34
35inputs = base_tokenizer(your_new_concept, return_tensors="pt").to("cuda")
36input_ids, attention_mask = inputs["input_ids"], inputs["attention_mask"]
37subspace_gen(input_ids, attention_mask)[0]
