Update download_model.py

download_model.py

@@ -1,10 +1,11 @@
 import torch
-import huggingface_hub
+import huggingface_hub import login
 from transformers import (AutoTokenizer, 
                           BitsAndBytesConfig, 
                           MBart50TokenizerFast,
                           AutoModelForSeq2SeqLM,
                           MBartForConditionalGeneration)
+from config import hf_token
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -26,10 +27,7 @@ def download_model(model_name: str):
         tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50", src_lang="en_XX", tgt_lang="en_XX")
         return model, tokenizer
     elif model_name == "Llama-3.2-1B-Instruct":
-        str1 = "f_bgSZT"
-        str2 = "AFSBqvApwHjMQuTOALqZKRpRBzEUL"
-        token = "h"+str1+str2
-        huggingface_hub.login(token = token)
+        login(token = hf_token)
         model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct").to(device)
         tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
         return model, tokenizer