Update app.py

app.py

@@ -7,11 +7,13 @@ class Net(nn.Module):
     def __init__(self):
         super(Net,self).__init__()
         self.layer = nn.Sequential(
-            nn.Linear(768, 768),
+            nn.Linear(768, 512),
             nn.ReLU(),
-            nn.Linear(768, 768),
+            nn.Linear(512, 256),
             nn.ReLU(),
-            nn.Linear(768, 8),
+            nn.Linear(256, 128),
+            nn.ReLU(),
+            nn.Linear(128, 8),
           )
 
     def forward(self,x):
@@ -32,18 +34,24 @@ def get_word_vector(sent, tokenizer, model):
  
      return get_hidden_states(encoded, model)
 
-model=Net()
-model.load_state_dict(torch.load('dummy_model.txt', map_location=torch.device('cpu')))
-model.eval()
 
 
 labels_articles = {1: 'Computer Science',2: 'Economics',3: "Electrical Engineering And Systems Science", 
 
                    4: "Mathematics",5: "Physics",6: "Quantitative Biology",7: "Quantitative Finance", 8: "Statistics"}
 
-tokenizer = AutoTokenizer.from_pretrained("Callidior/bert2bert-base-arxiv-titlegen")
-
-model_emb = AutoModelForSeq2SeqLM.from_pretrained("Callidior/bert2bert-base-arxiv-titlegen")
+@st.cache
+def load_models():
+    model=Net()
+    model.load_state_dict(torch.load('dummy_model.txt', map_location=torch.device('cpu')))
+    model.eval()
+    
+    tokenizer = AutoTokenizer.from_pretrained("Callidior/bert2bert-base-arxiv-titlegen")
+    
+    model_emb = AutoModelForSeq2SeqLM.from_pretrained("Callidior/bert2bert-base-arxiv-titlegen")
+    return model, model_emb, tokenizer
+
+model, model_emb, tokenizer = load_models()
 
 title = st.text_area("Write title of your article")
 summary = st.text_area("Write summary of your article or dont write anything (but you should press Ctrl + Enter)")