File size: 1,885 Bytes
b9662e5 |
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 |
---
license: cc-by-4.0
---
# Finetuned `Gemma-2-2B` for generating subspaces given any natural language descriptions for `Gemma-2-9B-it`
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!**
## How to use the subspace generator?
```py
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
class RegressionWrapper(torch.nn.Module):
def __init__(self, base_model, hidden_size, output_dim):
super().__init__()
self.base_model = base_model
self.regression_head = torch.nn.Linear(hidden_size, output_dim)
def forward(self, input_ids, attention_mask):
outputs = self.base_model.model(
input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True,
return_dict=True
)
last_hiddens = outputs.hidden_states[-1]
last_token_representations = last_hiddens[:, -1]
preds = self.regression_head(last_token_representations)
preds = F.normalize(preds, p=2, dim=-1)
return preds
base_model = AutoModelForCausalLM.from_pretrained(
f"google/gemma-2-2b", torch_dtype=torch.bfloat16)
base_tokenizer = AutoTokenizer.from_pretrained(
f"google/gemma-2-2b", model_max_length=512)
subspace_gen = RegressionWrapper(
base_model, hidden_size, output_dim).bfloat16().to("cuda")
subspace_gen.load_state_dict(torch.load('model.pth'))
your_new_concept = "terms related to Stanford University"
inputs = base_tokenizer(your_new_concept, return_tensors="pt").to("cuda")
input_ids, attention_mask = inputs["input_ids"], inputs["attention_mask"]
subspace_gen(input_ids, attention_mask)[0]
```
|