formatting, timing added

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import torch
 from torch import nn
 from transformers import BertModel, AutoTokenizer, AutoModel, pipeline
+from time import time
 # device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 device = 'cpu'
 
@@ -9,6 +10,7 @@ st.markdown("### Privet, mir!")
 st.markdown("<img width=200px src='https://i.pinimg.com/736x/11/33/19/113319f0ffe91f4bb0f468914b9916da.jpg'>", unsafe_allow_html=True)
 
 text = st.text_area("ENTER TEXT HERE")
+start_time = time()
 st.markdown("INFERENCE STARTS ...")
 
 # dict for decoding / enclding labels
@@ -34,20 +36,21 @@ class BertClassifier(nn.Module):
         linear_output = self.linear(dropout_output)
         final_layer = self.relu(linear_output)
         return final_layer
-        
-model = BertClassifier(n_classes=len(labels))
-st.markdown("Model created")
-model.load_state_dict(torch.load('model_weights_1.pt', map_location=torch.device('cpu')))
-model.eval()
-st.markdown("Model weights loaded")
 
+def build_model():
+    model = BertClassifier(n_classes=len(labels))
+    st.markdown("Model created")
+    model.load_state_dict(torch.load('model_weights_1.pt', map_location=torch.device('cpu')))
+    model.eval()
+    st.markdown("Model weights loaded")
+    return model
+    
+model = build_model()
+    
 def inference(txt, mode=None):
     # infers classes for text topic based on the trained model from above
     # has separate mode 'print' for just output
-    
-    txt = txt.lower().replace('\n', '')
-    
-    t2 = tokenizer(txt, 
+    t2 = tokenizer(txt.lower().replace('\n', ''), 
        padding='max_length', max_length = 512, truncation=True,
        return_tensors="pt")
 
@@ -58,19 +61,12 @@ def inference(txt, mode=None):
     out = out.cpu().detach().numpy().reshape(-1)
     out = out/out.sum() * 100
     res = [(l, o) for l, o in zip (list(labels.keys()), out.tolist())]   
-    
-    if mode == 'print':
-        res.sort(key = lambda x : - x[1])
-        for lbl, score in res:
-            if score >=1:
-                print(f"[{lbl:<7}] {labels_decoder[lbl]:<35}  {score:.1f}%")
-    
-    elif mode == 'debug':
-        return out, res
-    
-    else:
-        return res   
+    return res   
         
 res = inference(text, mode=None)    
 st.markdown("INFERENCE RESULT:")    
-st.markdown(f"{res}")
+for lbl, score in res:
+    if score >=1:
+        st.markdown(f"[{lbl:<7}] {labels_decoder[lbl]:<35}  {score:.1f}%")
+
+st.markdown(f"cycle time = {time() - start_time:.2f} s.")