Update models/attn_model.py

models/attn_model.py

@@ -1,3 +1,4 @@
+import os
 import torch
 from .model import Model
 from .utils import sample_token, get_last_attn
@@ -5,6 +6,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch.nn.functional as F
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
+token = os.getenv("HF_TOKEN") 
 
 class AttentionModel(Model):
     def __init__(self, config):
@@ -12,12 +14,14 @@ class AttentionModel(Model):
         self.name = config["model_info"]["name"]
         self.max_output_tokens = int(config["params"]["max_output_tokens"])
         model_id = config["model_info"]["model_id"]
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id,
+                            use_auth_token=token)
         self.model = AutoModelForCausalLM.from_pretrained(
             model_id,
             torch_dtype=torch.bfloat16,
             device_map=device,
-            attn_implementation="eager"
+            attn_implementation="eager",
+            use_auth_token=token
         ).eval()
         if config["params"]["important_heads"] == "all":
             attn_size = self.get_map_dim()