Update frontend.py

frontend.py

@@ -7,6 +7,24 @@ from model_utils import predict_proba
 st.set_page_config(page_title="StackOverflow Tagger", layout="wide")
 st.title("🔖 StackOverflow Tag Predictor")
 
+# ---- Choix du modèle dans la sidebar ----
+
+MODEL_OPTIONS = {
+    "BERT Overflow (maxcasado/BERT_overflow)": "maxcasado/BERT_overflow",
+    "Wendy Tags (wendyserver/predict_tags)": "wendyserver/predict_tags",
+}
+
+st.sidebar.header("⚙️ Configuration")
+model_label = st.sidebar.selectbox(
+    "Choisir le modèle",
+    list(MODEL_OPTIONS.keys()),
+)
+selected_model = MODEL_OPTIONS[model_label]
+
+st.sidebar.write(f"Modèle sélectionné : `{selected_model}`")
+
+# ---- Tabs : single question / CSV ----
+
 tab_single, tab_csv = st.tabs(["Question unique", "CSV batch"])
 
 with tab_single:
@@ -21,14 +39,24 @@ with tab_single:
         placeholder="Ex: How to fine-tune BERT for multi-label classification?",
     )
 
-    top_k = st.slider("Nombre de tags à afficher (top_k)", 1, 20, 5, key="topk_single")
+    top_k = st.slider(
+        "Nombre de tags à afficher (top_k)",
+        1,
+        20,
+        5,
+        key="topk_single",
+    )
 
     if st.button("Prédire", key="predict_single"):
         if not question.strip():
             st.warning("Merci d'entrer une question.")
         else:
-            with st.spinner("Prédiction en cours..."):
-                tags = predict_proba(question, top_k=top_k)
+            with st.spinner(f"Prédiction en cours avec {selected_model}..."):
+                tags = predict_proba(
+                    question,
+                    top_k=top_k,
+                    model_name=selected_model,
+                )
 
             if not tags:
                 st.warning("Pas de tags prédits.")
@@ -61,7 +89,7 @@ with tab_csv:
         )
 
         top_k_batch = st.slider(
-            "Nombre de tags à considérer pour le batch (pour choisir le meilleur)",
+            "Nombre de tags à considérer (pour choisir le meilleur)",
             1,
             20,
             5,
@@ -72,23 +100,33 @@ with tab_csv:
             if df[text_column].isnull().all():
                 st.error("La colonne choisie ne contient pas de texte.")
             else:
-                preds = []
-                with st.spinner("Prédiction en cours sur le CSV..."):
+                preds_best_tag = []
+                preds_best_score = []
+
+                with st.spinner(f"Prédiction batch avec {selected_model}..."):
                     for text in df[text_column].fillna(""):
-                        if not str(text).strip():
-                            preds.append({"best_tag": None, "best_score": None})
+                        s = str(text).strip()
+                        if not s:
+                            preds_best_tag.append(None)
+                            preds_best_score.append(None)
                             continue
-                        tags = predict_proba(str(text), top_k=top_k_batch)
+
+                        tags = predict_proba(
+                            s,
+                            top_k=top_k_batch,
+                            model_name=selected_model,
+                        )
+
                         if len(tags) == 0:
-                            preds.append({"best_tag": None, "best_score": None})
+                            preds_best_tag.append(None)
+                            preds_best_score.append(None)
                         else:
                             best = tags[0]
-                            preds.append(
-                                {"best_tag": best["label"], "best_score": best["score"]}
-                            )
+                            preds_best_tag.append(best["label"])
+                            preds_best_score.append(best["score"])
 
-                df["predicted_tag"] = [p["best_tag"] for p in preds]
-                df["predicted_score"] = [p["best_score"] for p in preds]
+                df["predicted_tag"] = preds_best_tag
+                df["predicted_score"] = preds_best_score
 
                 st.subheader("Résultats enrichis")
                 st.dataframe(df.head())