Changed code for loading model

app.py

@@ -1,13 +1,17 @@
 import gradio as gr
 import torch
 from src.model import BertClassifier, RobertaClassifier
+from transformers import BertTokenizer
+from datetime import datetime
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 try:
-    device_name = 'cuda' if torch.cuda.is_available() else 'cpu'
-    device = torch.device(device_name)
-    model = RobertaClassifier(model_name='roberta-base', dropout_rate=0.5)
-    model.to(device)
-    model.load_state_dict(torch.load('models/roberta-wiki.pth', map_location=device))
-    print('Model loaded')
+    model_name = 'bert-base-uncased'
+    model = BertClassifier(model_name, 0.5)
+    model.load_state_dict(torch.load('bert-all-data.pth', map_location=device))
+    tokenizer = BertTokenizer.from_pretrained(model_name)
+    print(f"[{datetime.now()}] Model loaded ")
+
 except:
     print('Model cannot be loaded')
 
@@ -18,4 +22,4 @@ def ai_text_classifier(text: str):
     }
 
 demo = gr.Interface(fn=ai_text_classifier, inputs="text", outputs="label")
-demo.launch(share=True)
+demo.launch()