OK

src/draft_model.py

@@ -2,6 +2,7 @@ from pathlib import Path
 import json
 import torch
 from typing import List, Dict
+from collections import defaultdict
 
 from huggingface_hub import hf_hub_download
 from helpers import get_embedding_dict, get_card_embeddings, MLP_CrossAttention
@@ -23,12 +24,12 @@ class DraftModel:
     def __init__(self):
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-        cfg_path = hf_hub_download(
-            repo_id=MODEL_REPO, filename=CFG_FILE, repo_type="model"
-        )
         weight_path = hf_hub_download(
             repo_id=MODEL_REPO, filename=MODEL_FILE, repo_type="model"
         )
+        cfg_path = hf_hub_download(
+            repo_id=MODEL_REPO, filename=CFG_FILE, repo_type="model"
+        )
 
         with open(cfg_path, "r") as f:
             cfg = json.load(f)
@@ -37,7 +38,7 @@ class DraftModel:
 
 
         self.net = MLP_CrossAttention(**cfg).to(self.device)
-        self.net.load_state_dict(weight_path, map_location=self.device)
+        self.net.load_state_dict(torch.load(weight_path, map_location=self.device))
         self.net.eval()
 
         # ---- embeddings – one-time load ------------------------------------
@@ -47,19 +48,43 @@ class DraftModel:
         )
         self.emb_size = next(iter(self.embed_dict.values())).shape[0]
 
+        raw_card_file = json.load(open(hf_hub_download(
+            repo_id=DATA_REPO, filename=CARD_FILE, repo_type="dataset"
+        )))
+        self.cards = defaultdict(dict)
+        for card in raw_card_file:
+            self.cards[card["set"]][card["name"]] = card
+        
+    
+    def _embed(self, name):  # helper
+        return get_card_embeddings((name,), embedding_dict=self.embed_dict)[0]
+
     # --------------------------------------------------------------------- #
     # Public API expected by streamlit_app.py                               #
     # --------------------------------------------------------------------- #
     @torch.no_grad()
-    def predict(self, pack: List[Dict], picks: List[Dict], deck: List[Dict]) -> Dict:
-        names = [c["name"] for c in pack]
-
-        def embed(name):  # helper
-            return get_card_embeddings((name,), embedding_dict=self.embed_dict)[0]
-
-        card_t = torch.stack([embed(n) for n in names]).unsqueeze(0).to(self.device)
-        deck_t = torch.zeros((1, 45, self.emb_size), device=self.device)
+    def predict(self, pack: List[Dict], deck: List[Dict]) -> Dict:
+
+        card_t = torch.stack([self._embed(c) for c in pack]).unsqueeze(0).to(self.device)
+        if deck is None:
+            deck_t = torch.zeros((1, 45, self.emb_size), device=self.device)
+        else:
+            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()
+        print(scores)
+        return {
+            "pick": pack[scores.argmax()],
+            "scores": scores.tolist(),
+        }
+    @torch.no_grad()
+    def get_p1p1(self, set_code:str):
+        keys = list(self.cards[set_code].keys())
+        cards = torch.stack([self._embed(c) for c in keys]).unsqueeze(0).to(self.device)
 
-        return torch.softmax(self.net(card_t, deck_t), dim=1).squeeze(0).cpu().numpy().tolist()
+        vals = self.predict(pack=keys, deck=None)["scores"]
+        return keys, vals
 
 

src/helper_files/supported_sets.txt

@@ -1,3 +1,5 @@
+EOE
+FIN
 DFT
 PIO
 FDN

src/helpers.py

@@ -280,4 +280,45 @@ def get_card_embeddings(card_names, embedding_dict, embedding_size = 1330):
         else:
             embedding, got_new = get_embedding_of_card(card, embedding_dict)
             embeddings.append(embedding)
-    return embeddings
+    return embeddings
+
+def check_for_basics(card_name, embedding_dict):
+    ints = ['1','2','3','4','5']
+    basics = ['Mountain','Forest','Swamp','Island','Plains']
+    for b in basics:
+        if b in card_name:
+            for i in ints:
+                if card_name == f'{b}_{i}':
+                    return b
+    return card_name
+    
+def get_embedding_of_card(card_name, embedding_dict):
+    try:
+        card_name = check_for_basics(card_name, embedding_dict)
+        card_name = card_name.replace('_', ' ')
+        card_name = card_name.replace("Sol'kanar","Sol'Kanar")
+        if card_name not in embedding_dict and card_name.split(' // ')[0] not in embedding_dict and card_name.replace('A-','') not in embedding_dict:
+            # print(f'Requesting new embedding for {card_name}')
+            # attributes, text = get_card_representation(card_name = card_name)
+            # text_embedding = embedd_text([text]).squeeze()
+            # return torch.Tensor(np.concatenate((attributes, text_embedding), axis = 0)), True
+            raise Exception(f'Could not find {card_name}')
+        else:
+            try:
+                return torch.Tensor(embedding_dict[card_name]), False
+            except:
+                try:
+                    return torch.Tensor(embedding_dict[card_name.split(' // ')[0]]), False
+                except:
+                    try:
+                        return torch.Tensor(embedding_dict[card_name.replace('_',' ')]), False
+                    except:
+                        try:
+                            return torch.Tensor(embedding_dict[card_name.replace('A-','')]), False
+                        except:
+                            print(f'Could not find {card_name}')
+                            raise Exception
+    except Exception as e:
+        print(f'Could not find {card_name}')
+        print(e)
+        raise e

src/streamlit_app.py

@@ -32,13 +32,14 @@ import random
 from pathlib import Path 
 
 from typing import Dict, List
+import pandas as pd
 
 import requests
 import streamlit as st
 
 from draft_model import DraftModel
 
-SUPPORTED_SETS_PATH = Path("supported_sets.txt")
+SUPPORTED_SETS_PATH = Path("src/helper_files/supported_sets.txt")
 
 
 @st.cache_data(show_spinner="Reading supported sets …")
@@ -54,6 +55,14 @@ def get_supported_sets(path: Path = SUPPORTED_SETS_PATH) -> List[str]:
 def load_model():
     return DraftModel() 
 
+@st.cache_data(show_spinner="Calculating P1P1 ...")
+def p1p1_ranking(set_code: str):
+    names, scores = model.get_p1p1(set_code)
+    df = pd.DataFrame({"card": names, "p1p1_score": scores})
+    df = df.sort_values("p1p1_score", ascending=False, ignore_index=True)
+    df.index += 1  # 1-based ranks look nicer
+    return df
+
 
 @st.cache_data(show_spinner="Calculating card rankings …")
 def rank_cards(set_code: str) -> List[Dict]:
@@ -82,11 +91,6 @@ def rank_cards(set_code: str) -> List[Dict]:
 
 model = load_model()
 
-
-# -----------------------------------------------------------------------------
-# 2. Draft‑logic helpers (stubs)
-# -----------------------------------------------------------------------------
-
 def suggest_pick(pack: List[Dict], picks: List[Dict]) -> Dict:
     if model is None:
         return random.choice(pack)
@@ -105,18 +109,6 @@ def fetch_card_image(card_name: str) -> str:
     return data["card_faces"][0]["image_uris"]["normal"]
 
 
-def generate_booster(set_code: str) -> List[Dict]:
-    url = f"https://api.scryfall.com/cards/search?q=set%3A{set_code}+is%3Abooster+unique%3Aprints"
-    cards: List[Dict] = []
-    while url:
-        resp = requests.get(url)
-        resp.raise_for_status()
-        payload = resp.json()
-        cards += payload["data"]
-        url = payload.get("next_page") if payload.get("has_more") else None
-    return random.sample(cards, 15)
-
-
 # -----------------------------------------------------------------------------
 # 3. Streamlit UI
 # -----------------------------------------------------------------------------
@@ -135,13 +127,19 @@ with st.sidebar:
     # Hide control in an expander (collapsed by default)
     with st.expander("Set selection", expanded=False):
         if supported_sets:
-            set_code = st.radio("Choose a set to draft", supported_sets, index=0)
+            # 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.",
+                "*supported_sets.txt* not found or empty. Using free-text input instead.",
                 icon="⚠️",
             )
-            set_code = st.text_input("Set code", value="WOE")
+            set_code = st.text_input("Set code", value="EOE", key="set_code")
 
         if st.button("Start new draft", type="primary"):
             st.session_state["pack"] = generate_booster(set_code)
@@ -154,7 +152,7 @@ st.session_state.setdefault("picks", [])
 
 # -------- Main content organised in tabs ------------------------------------
 
-tabs = st.tabs(["Draft", "Card rankings"])
+tabs = st.tabs(["Draft", "P1P1 Rankings"])
 
 # --- Tab 1: Draft ------------------------------------------------------------
 
@@ -189,11 +187,13 @@ with tabs[1]:
     st.header("Card rankings for set " + set_code)
 
     if set_code:
-        ranking = rank_cards(set_code)
-        if ranking:
-            for idx, card in enumerate(ranking, start=1):
-                st.write(f"{idx}. {card['name']} — **{card['score']:.2f}**")
-        else:
-            st.info("No cards found for this set (or Scryfall unavailable).")
+        try:
+            df = p1p1_ranking(set_code.lower())  # cached; auto-updates on dropdown change
+            if not df.empty:
+                st.dataframe(df, use_container_width=True)
+            else:
+                st.info("No P1P1 results returned for this set.")
+        except Exception as e:
+            st.error(f"Could not calculate P1P1: {e}")
     else:
-        st.info("Select a set in the sidebar to view rankings.")
+        st.info("Select a set in the sidebar to view P1P1.")