Update app.py

app.py

@@ -20,17 +20,20 @@ ID2CLS = {
 def classify(text, tokenizer, model):
     if not text:
         return [""]
-    tokens = tokenizer([text], truncation=True, padding=True, max_length=256, return_tensors="pt")['input_ids']
-    probabilities = torch.softmax(model(tokens).logits, dim=1).detach().cpu().numpy()[0]
+
+    batch = tokenizer([text], truncation=True, padding=True, max_length=256, return_tensors="pt")
+    outputs = model(batch['input_ids'], attention_mask=batch['attention_mask'])
+
+    probabilities = torch.softmax(outputs.logits, dim=1).detach().cpu().numpy()[0]
     total = 0
-    ans = []
 
     for p in probabilities.argsort()[::-1]:
-        if probabilities[p] + total < 0.9:
-            total += probabilities[p]
-            ans += [f'{ID2CLS[p]}: {round(probabilities[p] * 100, 2)}%']
+        field = f'{ID2CLS[p]}: {round(probabilities[p] * 100, 2)} %'
+        st.markdown(field)
 
-    return ans
+        total += probabilities[p]
+        if total > 0.95:
+            break
 
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
@@ -43,5 +46,4 @@ st.markdown("## Article classifier")
 title = st.text_area("title")
 text = st.text_area("article")
 
-for prediction in classify(title + text, tokenizer, model):
-    st.markdown(prediction)
+classify(title + text, tokenizer, model)