add probability feature

app.py

@@ -21,22 +21,31 @@ def lemmatize(s: str) -> iter:
     # lemmatize
     return map(lambda token: token.lemma_.lower(), tokens)
 
-def predict(title: str , post: str):
+def predict(title: str , post: str, predict_proba: bool):
     text = title + " " + post
     lemmes = np.array([' '.join(list(lemmatize(text)))])
 
     X = tfidf.transform(lemmes)
 
-    y_bin = model.predict(X)
-    y_tags = tags_binarizer.inverse_transform(y_bin)
+    if predict_proba:
+        y_proba = model.predict_proba(X)[0]
+        tags = list(dict(sorted(tags_binarizer.ts.count.items())).keys())
 
-    return y_tags
+        result = list(zip(tags, y_proba))
+    else:
+        y_bin = model.predict(X)
+        y_tags = tags_binarizer.inverse_transform(y_bin)
+
+        result = y_tags
+
+    return result
 
 demo = gr.Interface(
     fn=predict,
     inputs=[
-        gr.Textbox(lines=1, placeholder="Title..."),
-        gr.Textbox(lines=10, placeholder="Post...")],
+        gr.Textbox(label="Title", lines=1, placeholder="Title..."),
+        gr.Textbox(label="Post", lines=10, placeholder="Post..."),
+        gr.Checkbox(label="Proba?")],
     outputs=gr.Textbox(lines=10))
 
 demo.launch()