Update app.py

app.py

@@ -24,6 +24,21 @@ if tokenizer.pad_token is None:
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 gpt2.to(device)
 
+# Monkey patch the embed_to_distrib function to use the correct attribute
+def patched_embed_to_distrib(model, embed, log=True, logits=True):
+    if "gpt2" in model.config.architectures[0].lower():
+        with torch.inference_mode():
+            vocab = torch.matmul(embed, model.transformer.wte.weight.t())
+            if logits:
+                return vocab
+            if log:
+                return torch.log_softmax(vocab, dim=-1)
+            return torch.softmax(vocab, dim=-1)
+    else:
+        return pv.embed_to_distrib(model, embed, log, logits)
+
+pv.embed_to_distrib = patched_embed_to_distrib
+
 def simple_position_config(model_type, component, layer):
     config = IntervenableConfig(
         model_type=model_type,
@@ -118,13 +133,13 @@ async def process_inputs(base_sentence: str, rival_sentence: str):
 base_sentence = pn.widgets.TextInput(
     name="Base Sentence",
     placeholder="Enter the base sentence",
-    value="The cat is on the mat.",
+    value="Jane got some weird looks because she wore sunglasses outside at 4 PM.",
 )
 
 rival_sentence = pn.widgets.TextInput(
     name="Rival Sentence",
     placeholder="Enter the rival sentence",
-    value="The dog is on the mat.",
+    value="Jane got some weird looks because she wore sunglasses outside at 4 AM.",
 )
 
 input_widgets = pn.Column(