Update README.md

README.md

@@ -1,3 +1,49 @@
----
-license: cc-by-4.0
----
+---
+license: cc-by-4.0
+---
+
+# Finetuned `Gemma-2-2B` for generating subspaces given any natural language descriptions for `Gemma-2-2B-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]
+```