Display probabilities

app.py

@@ -44,13 +44,16 @@ def predict(model_abstract, model_claims, tokenizer_abstract, tokenizer_claims,
     attention_mask_claims = encoding_claims['attention_mask'].to(device)
     
     with torch.no_grad():
-        outputs_abstract = model_abstract(input_ids=input_abstract, attention_mask=attention_mask_abstract)
-        outputs_claims = model_claims(input_ids=input_claims, attention_mask=attention_mask_claims)
+        outputs_abstract = model_abstract(input_ids=input_abstract)
+        outputs_claims = model_claims(input_ids=input_claims)
+        
+    print(outputs_abstract.logits)
+    print(outputs_claims.logits)
         
     combined_prob = (outputs_abstract.logits.softmax(dim=1) + outputs_claims.logits.softmax(dim=1)) / 2
-    label = torch.argmax(combined_prob, axis=1).flatten()
+    label = torch.argmax(combined_prob, dim=1)
     
-    return label, combined_prob
+    return label, combined_prob.tolist()[0]
     
     
 if __name__ == '__main__':