Update app.py

app.py

@@ -1 +1,47 @@
-import torch
+import torch
+import numpy as np
+from transformers import AutoTokenizer, AutoModel
+import torch.nn as nn
+
+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"
+                  }
+
+class Net(nn.Module):
+    def __init__(self):
+        super(Net,self).__init__()
+        self.layer = nn.Sequential(
+            nn.Linear(768, 256),
+            nn.ReLU(),
+            nn.Linear(256, 128),
+            nn.ReLU(),
+            nn.Linear(128, 8),
+          )
+
+    def forward(self,x):
+        return self.layer(x)
+    
+model_second = Net()
+model_second.load_state_dict(torch.load('model.txt'))
+model_second.eval()
+
+model_name = 'bert-base-uncased'
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model_first = AutoModel.from_pretrained(model_name)
+
+title = st.text_area("Write the title of your article, please")
+abstract = st.text_area("Write the abstract")
+text = title + '. ' + abstract
+
+tokens_info = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
+out_first = model_first(**tokens_info).pooler_output
+out_second = model_second(out_first).detach().numpy()
+out_second = softmax(out_second)
+indices = np.argsort(out_second)[0][::-1]
+sum_prob = 0
+for i in indices:
+  print(labels_articles[i+1])
+  sum_prob += out_second[i]
+  if sum_prob > 0.95:
+    break