Update app.py

app.py

@@ -1,4 +1,4 @@
-from DistilBERT import model_DB
+from DistilBERT import DistilBERTClass
 import streamlit as st
 from transformers import DistilBertTokenizer, DistilBertModel
 import logging
@@ -9,6 +9,12 @@ MAX_LEN = 100
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased', truncation=True, do_lower_case=True)
 
+model_DB = DistilBERTClass()
+loaded_model_path = './model_DB_1.pt'
+model_DB.load_state_dict(torch.load(loaded_model_path, map_location=torch.device('cpu')))
+model_DB.to(device)
+
+
 def sentiment_analysis_DB(input):
     inputs = tokenizer.encode_plus(
         input,