Update app.py

app.py

@@ -1,4 +1,60 @@
 import streamlit as st
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+import torch.nn as nn
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+def get_hidden_states(encoded, model):
+     """Push input IDs through model. Stack and sum `layers` (last four by default).
+        Select only those subword token outputs that belong to our word of interest
+        and average them."""
+     with torch.no_grad():
+         output = model(decoder_input_ids=encoded['input_ids'], output_hidden_states=True, **encoded)
+
+     layers = [-4, -3, -2, -1]
+     states = output['decoder_hidden_states']
+     output = torch.stack([states[i] for i in layers]).sum(0).squeeze()
+ 
+     return output.mean(dim=0)
+ 
+def get_word_vector(sent, tokenizer, model):
+     encoded = tokenizer.encode_plus(sent, return_tensors="pt")
+ 
+     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")
+
+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)")
+
+text = title + '. ' + summary
+
+embed = get_word_vector(text, tokenizer, model_emb)
+
+logits = torch.nn.functional.softmax(model(embed), dim=0)
+
+best_tags = torch.argsort(logits, descending=True)
+
+sum = 0
+
+res = ''
+
+for tag in best_tags:
+  if sum > 0.95:
+    break
+  sum += logits[tag.item()]
+  # print(tag.item())
+  new_tag = labels_articles[tag.item() + 1]
+  res += new_tag + '\n'
+
+st.write('best tags = \n', res)