Update app.py

app.py

@@ -4,18 +4,21 @@ from src.model import BertClassifier, RobertaClassifier
 from transformers import BertTokenizer
 from datetime import datetime
 
+device = torch.device('cpu')
+model_name = 'bert-base-uncased'
+model = BertClassifier(model_name, 0.5)
+model.to(device)
+model.load_state_dict(torch.load('models/bert-all-data.pth', map_location=device))
+tokenizer = BertTokenizer.from_pretrained(model_name)
 
-
-def ai_text_classifier(text: str):
-    device = torch.device('cpu')
-    model_name = 'bert-base-uncased'
-    model = BertClassifier(model_name, 0.5)
-    model.to(device)
-    model.load_state_dict(torch.load('models/bert-all-data.pth', map_location=device))
-    tokenizer = BertTokenizer.from_pretrained(model_name)
+def ai_text_classifier(text: str):    
+    # Convert Text into tokens
     tokens = tokenizer(text, return_tensors='pt', max_length=512, padding='max_length', truncation=True).to(device)
+    
+    # Get probability of the text
     prob = model(tokens['input_ids'], tokens['attention_mask']).item()
     
+    # Return the probability in dictionary
     return {
         "AI": prob, 
         'Others': 1 - prob