modify inference for model file

inference.py

@@ -1,4 +1,3 @@
-import copy
 import re
 from pathlib import Path
 from typing import Generator
@@ -10,6 +9,7 @@ import tomotopy as tp  # type: ignore
 import torch
 import torch.nn as nn
 import transformers as T  # type: ignore
+from huggingface_hub import PyTorchModelHubMixin  # type: ignore
 from scipy import stats  # type: ignore
 from sudachipy import dictionary, tokenizer  # type: ignore
 
@@ -29,7 +29,7 @@ else:
     gpu = -1  # gpu = -1  # GPUが使用できなければ(CPUで処理)-1を指定
 
 
-cls_num = 3
+# cls_num = 3
 max_length = 512
 k_folds = 10
 bert_model_name = "cl-tohoku/bert-base-japanese-v3"
@@ -37,11 +37,10 @@ device = torch.device(f"cuda:{gpu}" if gpu>=0 else "cpu")
 
 
 #BERTモデルの定義
-class BertClassifier(nn.Module):
-    def __init__(self, model_name, cls_num=3):
-        super(BertClassifier, self).__init__()
-        #model_name = "cl-tohoku/bert-base-japanese"
-        self.bert = T.BertModel.from_pretrained(model_name, output_attentions=True)
+class BertClassifier(nn.Module, PyTorchModelHubMixin):
+    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)
 
         nn.init.normal_(self.fc.weight, std=0.02)
@@ -115,27 +114,31 @@ class SudachiTokenizer:
         return token_list
     
 
-def make_traind_model(bert_model):
+def make_traind_model():
     trained_models = []
     for k in range(k_folds):
         k = k + 1
-        model_path = model_base_path / f"trained_model{k}.pt"
-        trained_model = copy.deepcopy(bert_model)
-        trained_model.load_state_dict(torch.load(model_path, map_location=device), strict=False)
+        # model_path = model_base_path / f"trained_model{k}.pt"
+        # trained_model = copy.deepcopy(bert_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_models.append(trained_model)
     return trained_models
 
 
 @st.cache_resource
 def init_models():
-    bert_model = BertClassifier(bert_model_name, cls_num=1) #出力ノードを1に設定
-    bert_model.eval()
-    bert_model.to(device)
+    # bert_model = BertClassifier(cls_num=1) #出力ノードを1に設定
+    # bert_model.eval()
+    # bert_model.to(device)
 
     tokenizer_sudachi = SudachiTokenizer(split_mode="C")
     #Tokenizerの設定(ここではtokenizerをtokenizer_c2にしている)
     tokenizer_c2 = T.BertJapaneseTokenizer.from_pretrained(bert_model_name)
-    trained_models = make_traind_model(bert_model)
+    # trained_models = make_traind_model(bert_model)
+    trained_models = make_traind_model()
     return tokenizer_sudachi, tokenizer_c2, trained_models