Dzhamb commited on
Commit
ac09987
1 Parent(s): f8a5603

Update utils.py

Browse files
Files changed (1) hide show
  1. utils.py +3 -3
utils.py CHANGED
@@ -32,10 +32,10 @@ def get_labels(text, model, tokenizer, count_labels=8):
32
  return result_labels
33
  result_labels.append(pair[1])
34
 
35
-
36
  def load_model():
37
- tokenizer = DistilBertTokenizerFast()
38
- model = DistilBertForSequenceClassification()
39
  model.load_state_dict(torch.load('weight_model'))
40
  return model, tokenizer
41
 
 
32
  return result_labels
33
  result_labels.append(pair[1])
34
 
35
+ @st.cache(allow_output_mutation=True)
36
  def load_model():
37
+ tokenizer = DistilBertTokenizerFast
38
+ model = DistilBertForSequenceClassification
39
  model.load_state_dict(torch.load('weight_model'))
40
  return model, tokenizer
41