Works now

src/draft_model.py

@@ -74,10 +74,11 @@ class DraftModel:
             deck_t = torch.stack([self._embed(c) for c in deck]).unsqueeze(0).to(self.device)
 
         vals = self.net(deck = deck_t, cards = card_t)
-        #scores = torch.softmax(vals, dim=1).squeeze(0).cpu().numpy()
-        scores = vals.squeeze(0).cpu().numpy()
+        scores = torch.softmax(vals, dim=1).squeeze(0).cpu().numpy()
+        logits = vals.squeeze(0).cpu().numpy()
         return {
             "pick": pack[scores.argmax()],
+            "logits": logits.tolist(),
             "scores": scores.tolist(),
         }
     @torch.no_grad()
@@ -85,7 +86,7 @@ class DraftModel:
         keys = list(self.cards[set_code].keys())
         cards = torch.stack([self._embed(c) for c in keys]).unsqueeze(0).to(self.device)
 
-        vals = self.predict(pack=keys, deck=None)["scores"]
+        vals = self.predict(pack=keys, deck=None)["logits"]
         return keys, vals
 
 

src/streamlit_app.py

@@ -47,6 +47,39 @@ SUPPORTED_SETS_PATH = Path("src/helper_files/supported_sets.txt")
 def load_model():
     return DraftModel() 
 
+# --- callbacks ---
+def _add_selected_to_deck():
+    val = st.session_state.get("deck_selectbox")
+    if val:
+        add_card("deck", val)
+        st.session_state["deck_selectbox"] = None  # clear selection
+        st.toast(f"Added to deck: {val}")
+
+def _add_selected_to_pack():
+    val = st.session_state.get("pack_selectbox")
+    if val:
+        success = add_card("pack", val)
+        st.session_state["pack_selectbox"] = None  # clear selection
+        if success:
+            st.toast(f"Added to pack: {val}")
+        #st.rerun()
+
+def _reset_draft_state():
+    st.session_state["pack"] = []
+    st.session_state["deck"] = []
+    st.session_state["undo_stack"] = []
+    st.session_state["deck_selectbox"] = None
+    st.session_state["pack_selectbox"] = None
+
+def _on_set_changed():
+    curr = st.session_state.get("set_code")
+    prev = st.session_state.get("prev_set_code")
+    if prev != curr:
+        _reset_draft_state()
+        st.session_state["prev_set_code"] = curr
+        st.toast(f"Switched to set {curr}. Cleared current pack & deck.")
+
+
 
 if "model" not in st.session_state:
     st.session_state.model = load_model()  # your class
@@ -58,7 +91,7 @@ if "undo_stack" not in st.session_state:
     st.session_state.undo_stack: List[str] = []
 if "set_code" not in st.session_state:
     # choose a default set code that exists in model.cards, e.g., "eoe"
-    st.session_state.set_code = "eoe"
+    st.session_state.set_code = "EOE"
 
 model = st.session_state.model
 
@@ -88,8 +121,9 @@ def rank_cards(deck: List[str], pack: List[str]) -> List[Dict]:
     else:
         out = model.predict(pack, deck = deck)
     pick = out["pick"]
+    logits = {pack[i]: score for i, score in enumerate(out["logits"])}
     scores = {pack[i]: score for i, score in enumerate(out["scores"])}
-    return pick, scores
+    return pick, logits, scores
 
 
 def fetch_card_image(card_name: str) -> str:
@@ -118,29 +152,13 @@ with st.sidebar:
     st.header("Draft setup")
 
     supported_sets = get_supported_sets()
-
-    # Hide control in an expander (collapsed by default)
-    if supported_sets:
-        # UPDATED: dropdown instead of radio
-        set_code = st.selectbox(
-            "Choose a set",
-            supported_sets,
-            index=0,
-            key="set_code",
-        )
-    else:
-        st.warning(
-            "*supported_sets.txt* not found or empty. Using free-text input instead.",
-            icon="⚠️",
-        )
-        set_code = st.text_input("Set code", value="EOE", key="set_code")
-
-
-# -------- Session state ------------------------------------------------------
-
-st.session_state.setdefault("pack", [])
-st.session_state.setdefault("picks", [])
-st.session_state.setdefault("undo_stack", [])
+    set_code = st.selectbox(
+        "Choose a set",
+        supported_sets,
+        index=0,
+        key="set_code",
+        on_change=_on_set_changed
+    )
 
 # -------- Main content organised in tabs ------------------------------------
 
@@ -148,7 +166,15 @@ tabs = st.tabs(["Draft", "P1P1 Rankings"])
 
 def add_card(target: str, card: str):
     """target is 'pack' or 'picks'."""
-    st.session_state[target].append(card)
+    if target == "pack":
+        if card not in st.session_state["pack"]:
+            st.session_state[target].append(card)
+        else:
+            st.warning(f"{card} is already in the pack.", icon="⚠️")
+            return False
+    elif target == "deck":
+        st.session_state[target].append(card)
+    return True
 
 def remove_card(target: str, key: str):
     lst = st.session_state[target]
@@ -160,7 +186,7 @@ def push_undo():
     """Save a snapshot of pack + picks so we can undo one step."""
     st.session_state["undo_stack"].append({
         "pack": copy.deepcopy(st.session_state["pack"]),
-        "picks": copy.deepcopy(st.session_state["picks"]),
+        "deck": copy.deepcopy(st.session_state["deck"]),
     })
     # (optional) cap history
     if len(st.session_state["undo_stack"]) > 20:
@@ -170,29 +196,13 @@ def undo_last():
     if st.session_state.get("undo_stack"):
         snap = st.session_state["undo_stack"].pop()
         st.session_state["pack"]  = snap["pack"]
-        st.session_state["picks"] = snap["picks"]
+        st.session_state["deck"] = snap["deck"]
 
-# --- callbacks ---
-def _add_selected_to_deck():
-    val = st.session_state.get("deck_selectbox")
-    if val:
-        add_card("picks", val)
-        st.session_state["deck_selectbox"] = None  # clear selection
-        st.toast(f"Added to deck: {val}")
-
-def _add_selected_to_pack():
-    val = st.session_state.get("pack_selectbox")
-    if val:
-        add_card("pack", val)
-        st.session_state["pack_selectbox"] = None  # clear selection
-        st.toast(f"Added to pack: {val}")
-        #st.rerun()
 
 
-# --- Tab 1: Draft ------------------------------------------------------------
-
+# --- Tab 1: Draft -------------------------------------------------------
 with tabs[0]:
-    
+
     if st.session_state["undo_stack"]:
         st.button("↩️ Undo last action", on_click=undo_last)
 
@@ -201,78 +211,95 @@ with tabs[0]:
 
     if st.session_state["pack"]:
         pack = st.session_state["pack"]
-        deck = st.session_state["picks"]
+        deck = st.session_state["deck"]
 
         try:
-            pick, scores = rank_cards(deck, pack)
+            pick, logits, scores = rank_cards(deck, pack)
+            if pick:
+                st.success(f"💡 Suggested pick: **{pick}**", icon="✨")
         except Exception as e:
             st.error(f"Error calculating card rankings: {e}")
-    if not set_code:
-        st.info("Pick a set to draft")
-    else:
-        options = list(model.cards[set_code.lower()].keys())
-        c1, c2 = st.columns(2)
-        with c1:
-            st.subheader("Add to Deck")
-            deck_sel = st.selectbox(
-                "Search card (deck)",
-                options,
-                index=None,
-                placeholder="Type to search…",
-                key="deck_selectbox",
-                on_change=_add_selected_to_deck,   # <- auto-add
-            )
-
-            # Show current deck with remove buttons
-            if st.session_state["picks"]:
-                for i, card in enumerate(st.session_state["picks"]):
-                    name_col, rm_col = st.columns([6, 3], gap="small")
-                    name_col.write(card)
-                    with rm_col:
-                        if st.button("Remove", key=f"rm-deck-{i}", use_container_width=True):
-                            remove_card("picks", card)
-                            st.rerun()
-            else:
-                st.caption("Deck is empty.")
-
-        with c2:
-            st.subheader("Add to pack")
-            pack_sel = st.selectbox(
-                "Search card (pack)",
-                options,
-                index=None,
-                placeholder="Type to search…",
-                key="pack_selectbox",
-                on_change=_add_selected_to_pack,   # <- auto-add
-            )
+
+    options = list(model.cards[set_code.lower()].keys())
+    c1, c2 = st.columns(2)
+    with c1:
+        st.subheader("Add to Deck")
+        deck_sel = st.selectbox(
+            "Search card (deck)",
+            options,
+            index=None,
+            placeholder="Type to search…",
+            key="deck_selectbox",
+            on_change=_add_selected_to_deck,   # <- auto-add
+        )
+        if st.session_state["deck"]:
+            # header row
+            st.button("🗑️ Clear deck", on_click=lambda: st.session_state.update(deck=[]), use_container_width=True)
+            h1, h2 = st.columns([6, 3])
+            h1.markdown("**Card**")
+            h2.markdown("**Remove?**")
+            for i, card in enumerate(st.session_state["deck"]):
+                name_col, rm_col = st.columns([6, 3], gap="small")
+                name_col.write(card)
+                with rm_col:
+                    if st.button("Remove", key=f"rm-deck-{i}", use_container_width=True):
+                        remove_card("deck", card)
+                        st.rerun()
+        else:
+            st.caption("Deck is empty.")
+
+    with c2:
+        st.subheader("Add to pack")
+        pack_sel = st.selectbox(
+            "Search card (pack)",
+            options,
+            index=None,
+            placeholder="Type to search…",
+            key="pack_selectbox",
+            on_change=_add_selected_to_pack,   # <- auto-add
+        )
+            
+        
+
+        if st.session_state["pack"]:
+            # header row
+            st.button("🗑️ Clear pack", on_click=lambda: st.session_state.update(pack=[]), use_container_width=True)
+            h1, h2, h3 = st.columns([6, 2, 3])
+            h1.markdown("**Card**")
+            h2.markdown("**Score**")
+            h3.markdown("**Pick?**")
+
+            pack_list = st.session_state["pack"]
+            vals = [scores.get(c) if scores and c in scores else np.nan for c in pack_list]
+            logits = [logits.get(c) if logits and c in logits else np.nan for c in pack_list]
+            df_scores = pd.DataFrame({"Card": pack_list, "Score": vals, "Logits": logits})
+            df_scores = df_scores.sort_values("Score", ascending=False, na_position="last").reset_index(drop=True)
+            
+            
+            # rows
+            for i, row in df_scores.iterrows():
+                card = row["Card"]
+                score = row["Score"]
+                logit = row["Logits"]
+
+
+                c1, c2, c3 = st.columns([6, 2, 3], gap="small")
+                c1.write(card)
                 
-            if st.session_state["pack"]:
-                pack_list = st.session_state["pack"]
-                vals = [scores.get(c) if scores and c in scores else np.nan for c in pack_list]
-                df_scores = pd.DataFrame({"Card": pack_list, "Score": vals})
-
-                pack_list = st.session_state["pack"]
-                vals = [scores.get(c) if scores and c in scores else np.nan for c in pack_list]
-
-                # header row
-                h1, h2, h3 = st.columns([6, 2, 3])
-                h1.markdown("**Card**")
-                h2.markdown("**Score**")
-                h3.markdown("**Pick**")
-
-
-                # rows
-                for i, card in enumerate(pack_list):
-                    s = vals[i]
-                    c1, c2, c3 = st.columns([6, 2, 3], gap="small")
-                    c1.write(card)
-                    c2.markdown(f"<div class='score-cell'>{'' if np.isnan(s) else f'{s:.4f}'}</div>", unsafe_allow_html=True)
-                    with c3:
-                        if st.button("Pick", key=f"pick_btn_{i}", use_container_width=True, help="Add to deck & clear pack"):
-                            push_undo()
-                            st.session_state["picks"].append(card)
-                            st.session_state["pack"] = []   # or generate_booster(set_code)
-                            st.rerun()
+                tooltip_html = f"""
+                <div title="{logit:.4f}">
+                    <progress Value="{score}" max="1" style="width: 100%; height: 20px;"></progress>
+                </div>
+                """
+                c2.markdown(tooltip_html, unsafe_allow_html=True)
+                with c3:
+                    if st.button("Pick", key=f"pick_btn_{i}", use_container_width=True, help="Add to deck & clear pack"):
+                        push_undo()
+                        st.session_state["deck"].append(card)
+                        st.session_state["pack"] = [] 
+                        st.rerun()
+        else:
+            st.caption("Pack is empty.")
 
 # --- Tab 2: Card rankings ----------------------------------------------------
 
@@ -289,4 +316,4 @@ with tabs[1]:
         except Exception as e:
             st.error(f"Could not calculate P1P1: {e}")
     else:
-        st.info("Select a set in the sidebar to view P1P1.")
+        st.info("Select a set in the sidebar to view P1P1.")