Create model.py

model.py

@@ -0,0 +1,70 @@
+# model.py
+self.model = GPT2LMHeadModel.from_pretrained(model_name, output_hidden_states=True).to(self.device)
+self.model.eval()
+
+
+def generate_text(self, prompt, max_length=50, top_k=10):
+inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+with torch.no_grad():
+output = self.model.generate(**inputs, max_length=len(inputs['input_ids'][0]) + max_length, do_sample=True, top_k=top_k, pad_token_id=self.tokenizer.eos_token_id)
+return self.tokenizer.decode(output[0], skip_special_tokens=True)
+
+
+def _get_hidden_states(self, prompt):
+inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+with torch.no_grad():
+out = self.model(**inputs)
+# hidden_states: tuple(len = n_layers+1) of (batch, seq_len, hidden)
+return out.hidden_states
+
+
+def layer_importance(self, prompt, experiment_type="story_continuation"):
+"""
+Simple proxy for activation patching: measure how sensitive the model's next-token logits are
+to zeroing the output of each transformer block (layer). For each layer:
+- compute logits with all layers active
+- compute logits with layer `l` zeroed out (set its hidden output to zero)
+- compute L1 difference between the top token logits — larger difference => higher importance
+
+
+Returns a list of importance scores (one per transformer block).
+"""
+# 1) get baseline logits for the prompt
+inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
+input_ids = inputs['input_ids']
+with torch.no_grad():
+out = self.model(**inputs, output_hidden_states=True)
+baseline_logits = out.logits[0, -1, :].cpu().numpy()
+
+
+# Number of transformer blocks
+n_layers = len(out.hidden_states) - 1
+scores = []
+
+
+# We'll re-run forward passes while zeroing each layer's output using a forward hook
+for layer_idx in range(n_layers):
+def hook(module, inp, outp):
+# outp has shape (batch, seq_len, hidden)
+return torch.zeros_like(outp)
+
+
+# register hook on transformer.h.{layer_idx}
+handle = self.model.transformer.h[layer_idx].register_forward_hook(hook)
+with torch.no_grad():
+out2 = self.model(**inputs)
+logits2 = out2.logits[0, -1, :].cpu().numpy()
+diff = np.sum(np.abs(baseline_logits - logits2))
+scores.append(float(diff))
+handle.remove()
+
+
+# Normalize scores to 0-1
+arr = np.array(scores)
+if arr.max() > 0:
+arr = (arr - arr.min()) / (arr.max() - arr.min())
+else:
+arr = arr * 0.0
+
+
+return arr.tolist()