Dzhamb commited on
Commit
99ca19b
1 Parent(s): ac09987

Update utils.py

Browse files
Files changed (1) hide show
  1. utils.py +2 -2
utils.py CHANGED
@@ -34,8 +34,8 @@ def get_labels(text, model, tokenizer, count_labels=8):
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
 
 
34
 
35
  @st.cache(allow_output_mutation=True)
36
  def load_model():
37
+ tokenizer = DistilBertTokenizerFast.from_pretrained("distilbert-base-cased")
38
+ model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-cased", num_labels=8)
39
  model.load_state_dict(torch.load('weight_model'))
40
  return model, tokenizer
41