Update app.py

app.py

@@ -3,6 +3,70 @@ import torch
 import torch.nn as nn
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
+class Net(nn.Module):
+    def __init__(self):
+        super(Net,self).__init__()
+        self.layer = nn.Sequential(
+            nn.Linear(768, 512),
+            nn.ReLU(),
+            nn.Linear(512, 256),
+            nn.ReLU(),
+            nn.Linear(256, 128),
+            nn.ReLU(),
+            nn.Linear(128, 8),
+          )
+
+    def forward(self,x):
+        return self.layer(x)
+    
+model = Net()
+model.load_state_dict(torch.load('model.dat'))
+tokenizer = AutoTokenizer.from_pretrained("Callidior/bert2bert-base-arxiv-titlegen")
+model_emb = AutoModelForSeq2SeqLM.from_pretrained("Callidior/bert2bert-base-arxiv-titlegen")
+
+def BuildAnswer(txt):
+    def get_hidden_states(encoded, model):
+        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",  truncation=True)
+        return get_hidden_states(encoded, model)
+        
+    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"
+    }
+    
+    embed = get_word_vector(txt, tokenizer, model_emb)
+    logits = torch.nn.functional.softmax(model(embed), dim=0)
+    best_tags = torch.argsort(logits, descending=True)
+    
+    sum = 0
+    result = []
+    for tag in best_tags:
+        if sum > 0.95:
+            break
+        sum += logits[tag.item()]
+        res = round(float(logits[tag.item()].cpu()) * 100)
+        label = labels_articles[tag.item() + 1]
+        result.append(f'{res:3d}% - {label}')
+    return '\n'.join(result)
+
+
+
 st.markdown("### Hello, world!")
 st.markdown("<img width=200px src='https://rozetked.me/images/uploads/dwoilp3BVjlE.jpg'>", unsafe_allow_html=True)
 # ^-- можно показывать пользователю текст, картинки, ограниченное подмножество html - всё как в jupyter
@@ -17,4 +81,4 @@ abstract = st.text_area("Abstract")
 #raw_predictions = pipe(text)
 # тут уже знакомый вам код с huggingface.transformers -- его можно заменить на что угодно от fairseq до catboost
 
-st.markdown(f"{title + ' ' + abstract}")
+st.markdown(f"{BuildAnswer(title + ' ' + abstract)}")