add returning probabilities as well

app.py

@@ -14,9 +14,16 @@ def pipeline(video_url):
     video_id = video_url.split("?v=")[-1]
     punctuated_text = punctuate(video_id)
     sentences = re.split(r"[\.\!\?]\s", punctuated_text)
-    classification = predict_from_document(sentences)
+    classification, probs = predict_from_document(sentences)
     #    return punctuated_text
-    return [{"start": "12:05", "end": "12:52", "classification": str(classification)}]
+    return [
+        {
+            "start": "12:05",
+            "end": "12:52",
+            "classification": str(classification),
+            "probabilities": probs,
+        }
+    ]
 
 
 # print(pipeline("VL5M5ZihJK4"))

functions/model_infer.py

@@ -38,4 +38,4 @@ def predict_from_document(sentences):
     preprop = preprocess(sentences)
     prediction = model.predict(preprop)
     output = (prediction.flatten()[: len(sentences)] >= 0.5).astype(int)
-    return output
+    return output, prediction.flatten()[: len(sentences)]