Update app.py

app.py

@@ -6,23 +6,10 @@ import streamlit as st
 from transformers import DistilBertModel
 from tokenization_kobert import KoBertTokenizer
 
-@st.cache(allow_output_mutation=True)
-def get_model():
-    bert_model = DistilBertModel.from_pretrained('monologg/distilkobert')
-    tokenizer = KoBertTokenizer.from_pretrained('monologg/distilkobert')
-    
-    model = SanctiMoly(freeze_bert=False)
-    checkpoint = torch.load("./model.pt", map_location=device)
-    model.load_state_dict(checkpoint['model_state_dict'])
-    
-    return model, tokenizer
-    
-model, tokenizer = get_model()
-
 class SanctiMoly(nn.Module):
     """ Holy Moly News BERT """
 
-    def __init__(self, freeze_bert = True):
+    def __init__(self, bert_model, freeze_bert = True):
         super(SanctiMoly, self).__init__()
         self.encoder = bert_model
         # FC-BN-Tanh
@@ -57,6 +44,21 @@ class SanctiMoly(nn.Module):
         # print(output.shape)
         return output
 
+@st.cache(allow_output_mutation=True)
+def get_model():
+    bert_model = DistilBertModel.from_pretrained('monologg/distilkobert')
+    tokenizer = KoBertTokenizer.from_pretrained('monologg/distilkobert')
+    
+    model = SanctiMoly(bert_model, freeze_bert=False)
+    checkpoint = torch.load("./model.pt", map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    
+    return model, tokenizer
+    
+model, tokenizer = get_model()
+
+
+
 
 class RegexSubstitution(object):
     """Regex substitution class for transform"""