DivEye - PR (token) (#4)

- token changes (0392639a5f17fa83c58b523b8c26a1556f6da092)

software.py

@@ -10,6 +10,7 @@ from tqdm import tqdm
 from torch.nn import CrossEntropyLoss
 from pathlib import Path
 import spaces
+import os
 
 class Diversity:
     def __init__(self, model, tokenizer, device):
@@ -103,14 +104,16 @@ class Software:
         else:
             self.device_bi = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-        self.div_tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b", use_fast=False, trust_remote_code=True)
+        self.token = os.getenv("HF_TOKEN") 
+
+        self.div_tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b", use_fast=False, trust_remote_code=True, use_auth_token=self.token)
         self.div_model = AutoModelForCausalLM.from_pretrained(
-            "tiiuae/falcon-7b", device_map=self.device_div, torch_dtype=torch.float16, trust_remote_code=True
+            "tiiuae/falcon-7b", device_map=self.device_div, torch_dtype=torch.float16, trust_remote_code=True, use_auth_token=self.token
         )
-        
-        self.bi_tokenizer = AutoTokenizer.from_pretrained("google/gemma-1.1-2b-it", use_fast=False, trust_remote_code=True)
+
+        self.bi_tokenizer = AutoTokenizer.from_pretrained("google/gemma-1.1-2b-it", use_fast=False, trust_remote_code=True, use_auth_token=self.token)
         self.bi_model = AutoModelForCausalLM.from_pretrained(
-            "google/gemma-1.1-2b-it", device_map=self.device_bi, torch_dtype=torch.float16, trust_remote_code=True
+            "google/gemma-1.1-2b-it", device_map=self.device_bi, torch_dtype=torch.float16, trust_remote_code=True, use_auth_token=self.token
         )
         
         self.diveye = Diversity(self.div_model, self.div_tokenizer, self.device_div)