modify get model with token

inference.py

@@ -1,3 +1,4 @@
+import os
 import re
 from pathlib import Path
 from typing import Generator
@@ -13,7 +14,7 @@ from huggingface_hub import PyTorchModelHubMixin  # type: ignore
 from scipy import stats  # type: ignore
 from sudachipy import dictionary, tokenizer  # type: ignore
 
-# from transformers import AutoModel
+HF_AUTH_TOKEN = os.getenv("HF_AUTH_TOKEN")
 
 MODELS_PATH = Path(__file__).parent / "saved_model"
 # model_base_path = MODELS_PATH / "two_class"
@@ -38,7 +39,7 @@ device = torch.device(f"cuda:{gpu}" if gpu>=0 else "cpu")
 
 #BERTモデルの定義
 class BertClassifier(nn.Module, PyTorchModelHubMixin):
-    def __init__(self, cls_num: int = 1):
+    def __init__(self, cls_num: int):
         super().__init__()
         self.bert = T.BertModel.from_pretrained(bert_model_name, output_attentions=True)
         self.fc = nn.Linear(768, cls_num, bias=True)
@@ -123,7 +124,7 @@ def make_traind_model():
         # trained_model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
         # trained_models.append(trained_model)
         model_name = MODEL_BASE + str(k)
-        trained_model = BertClassifier.from_pretrained(model_name).to(device)
+        trained_model = BertClassifier.from_pretrained(model_name, token=HF_AUTH_TOKEN).to(device)
         print(f"Got model {model_name}")
         trained_models.append(trained_model)
     return trained_models