logging in

model.py

@@ -1,21 +1,32 @@
+import os
 from threading import Thread
 from typing import Iterator
 
 import torch
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from huggingface_hub import login
 
 # model_id = 'meta-llama/Llama-2-13b-chat-hf'
 model_id = 'meta-llama/Llama-2-7b-chat-hf'
 
 if not torch.cuda.is_available():
-    config = AutoConfig.from_pretrained(model_id)
+    tok = os.environ['HF_TOKEN']
+    login(new_session=True,
+          write_permission=False,
+          token=tok
+          
+          #, token="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"
+          )
+    
+    config = AutoConfig.from_pretrained(model_id, use_auth_token=True)
     config.pretraining_tp = 1
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         config=config,
         torch_dtype=torch.float16,
         load_in_4bit=True,
-        device_map='auto'
+        device_map='auto',
+        use_auth_token=True
     )
 else:
     model = None