Works now

src/draft_model.py

@@ -22,7 +22,8 @@ ENCODING_FILE = "card_encodings.pt"
 
 class DraftModel:
     def __init__(self):
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        #self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.device = torch.device("cpu")  # Force CPU for compatibility
 
         weight_path = hf_hub_download(
             repo_id=MODEL_REPO, filename=MODEL_FILE, repo_type="model"
@@ -37,9 +38,10 @@ class DraftModel:
         # ---- load network ---------------------------------------------------
 
 
-        self.net = MLP_CrossAttention(**cfg).to(self.device)
+        self.net = MLP_CrossAttention(**cfg)
         self.net.load_state_dict(torch.load(weight_path, map_location=self.device))
         self.net.eval()
+        self.net.to(self.device)
 
         # ---- embeddings – one-time load ------------------------------------
         self.embed_dict = get_embedding_dict(

src/helpers.py

@@ -194,13 +194,13 @@ class MLP_CrossAttention(nn.Module):
         
         deck = deck.view(batch_size * deck_size, card_size)
 
-        deck_encoded = self.card_encoder(deck.cuda())
+        deck_encoded = self.card_encoder(deck)
         deck_encoded = deck_encoded.view(batch_size, deck_size, -1)
 
 
         # identify padded cards
-        mask = (cards.sum(dim=-1) != 0).cuda()
-        cards_encoded = self.card_encoder(cards.cuda())
+        mask = (cards.sum(dim=-1) != 0)
+        cards_encoded = self.card_encoder(cards)
 
         if not no_attention:
             # Cross-attention

src/streamlit_app.py

@@ -47,7 +47,22 @@ SUPPORTED_SETS_PATH = Path("src/helper_files/supported_sets.txt")
 def load_model():
     return DraftModel() 
 
-model = load_model()
+
+if "model" not in st.session_state:
+    st.session_state.model = load_model()  # your class
+if "deck" not in st.session_state:
+    st.session_state.deck: List[str] = []
+if "pack" not in st.session_state:
+    st.session_state.pack: List[str] = []
+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"
+
+model = st.session_state.model
+
+
 
 @st.cache_data(show_spinner="Reading supported sets …")
 def get_supported_sets(path: Path = SUPPORTED_SETS_PATH) -> List[str]:
@@ -68,7 +83,10 @@ def p1p1_ranking(set_code: str):
 
 @st.cache_data(show_spinner="Calculating card rankings …")
 def rank_cards(deck: List[str], pack: List[str]) -> List[Dict]:
-    out = model.predict(pack, deck)
+    if not deck:
+        out = model.predict(pack, deck = None)
+    else:
+        out = model.predict(pack, deck = deck)
     pick = out["pick"]
     scores = {pack[i]: score for i, score in enumerate(out["scores"])}
     return pick, scores
@@ -127,19 +145,12 @@ st.session_state.setdefault("undo_stack", [])
 # -------- Main content organised in tabs ------------------------------------
 
 tabs = st.tabs(["Draft", "P1P1 Rankings"])
-def ensure_state():
-    """Make sure session_state has the pack and picks lists."""
-    if "pack" not in st.session_state:
-        st.session_state["pack"] = []
-    if "picks" not in st.session_state:
-        st.session_state["picks"] = []
+
 def add_card(target: str, card: str):
     """target is 'pack' or 'picks'."""
-    ensure_state()
     st.session_state[target].append(card)
 
 def remove_card(target: str, key: str):
-    ensure_state()
     lst = st.session_state[target]
     idx = next((i for i, c in enumerate(lst) if c == key), None)
     if idx is not None:
@@ -147,7 +158,6 @@ def remove_card(target: str, key: str):
 
 def push_undo():
     """Save a snapshot of pack + picks so we can undo one step."""
-    ensure_state()
     st.session_state["undo_stack"].append({
         "pack": copy.deepcopy(st.session_state["pack"]),
         "picks": copy.deepcopy(st.session_state["picks"]),
@@ -161,7 +171,6 @@ def undo_last():
         snap = st.session_state["undo_stack"].pop()
         st.session_state["pack"]  = snap["pack"]
         st.session_state["picks"] = snap["picks"]
-        st.rerun()
 
 # --- callbacks ---
 def _add_selected_to_deck():
@@ -170,7 +179,6 @@ def _add_selected_to_deck():
         add_card("picks", val)
         st.session_state["deck_selectbox"] = None  # clear selection
         st.toast(f"Added to deck: {val}")
-        st.rerun()
 
 def _add_selected_to_pack():
     val = st.session_state.get("pack_selectbox")
@@ -178,14 +186,12 @@ def _add_selected_to_pack():
         add_card("pack", val)
         st.session_state["pack_selectbox"] = None  # clear selection
         st.toast(f"Added to pack: {val}")
-        st.rerun()
+        #st.rerun()
 
 
 # --- Tab 1: Draft ------------------------------------------------------------
 
 with tabs[0]:
-    ensure_state()
-
     
     if st.session_state["undo_stack"]:
         st.button("↩️ Undo last action", on_click=undo_last)
@@ -220,11 +226,12 @@ with tabs[0]:
             # Show current deck with remove buttons
             if st.session_state["picks"]:
                 for i, card in enumerate(st.session_state["picks"]):
-                    cols = st.columns([6, 2])
-                    cols[0].write(card)
-                    if cols[1].button("Remove", key=f"rm-deck-{i}"):
-                        remove_card("picks", card)
-                        st.rerun()
+                    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.")
 
@@ -243,26 +250,29 @@ with tabs[0]:
                 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})
-                tbl_col, btn_col = st.columns([4, 1], gap="small")
-                with tbl_col:
-                    st.dataframe(
-                        df_scores,
-                        use_container_width=True,
-                        column_config={
-                            "Score": st.column_config.NumberColumn(format="%.4f")
-                        },
-                        hide_index=True,
-                    )
-                with btn_col:
-                    for i, card in enumerate(pack_list):
-                        cols = st.columns([8, 4])
-                        if cols[1].button("Add", key=f"add_clear_{i}"):
+
+                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.session_state["pack"] = []   # or generate_booster(set_code)
                             st.rerun()
-            else:
-                st.caption("Pack is empty.")
 
 # --- Tab 2: Card rankings ----------------------------------------------------