changed output logic

app.py

@@ -55,10 +55,12 @@ def predict(text, model_id, tokenizer_id):
         logits = model(**inputs).logits
 
     probs = torch.nn.functional.softmax(logits, dim=1).cpu().numpy().flatten()
-    label_pred = model.config.id2label[probs.argmax()]
-    probability_pred = f"{round(100*probs.max(), 2)}%"
-    return label_pred, probability_pred
+    return probs
 
+def get_most_probable_label(probs):
+    label = model.config.id2label[probs.argmax()]
+    probability = f"{round(100 * probs.max(), 2)}%"
+    return [sentence, label, probability]
 
 def predict_wrapper(text, language):
     model_id = build_huggingface_path(language)
@@ -69,8 +71,8 @@ def predict_wrapper(text, language):
 
     results = []
     for sentence in sentences:
-        label, probability = predict(sentence, model_id, tokenizer_id)
-        results.append([sentence, label, probability])
+        probs = predict(sentence, model_id, tokenizer_id)
+        results.append(get_most_probable_label(probs))
 
     output_info = f'Prediction was made using the <a href="https://huggingface.co/{model_id}">{model_id}</a> model.'
     return results, output_info