Spaces:
Sleeping
Sleeping
| 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 | |
| MODEL_REPO = "TimoBertram/MTG_Model_FIN" | |
| CFG_FILE = "config.json" | |
| MODEL_FILE = "network.pt" | |
| DATA_REPO = "TimoBertram/MTG_Drafting_Dataset" | |
| CARD_FILE = "cards_eoe.json" | |
| 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("cpu") # Force CPU for compatibility | |
| 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) | |
| # ---- load network --------------------------------------------------- | |
| 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( | |
| hf_hub_download(repo_id=DATA_REPO, filename=ENCODING_FILE, repo_type="dataset"), | |
| add_nontransformed=True | |
| ) | |
| 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 # | |
| # --------------------------------------------------------------------- # | |
| def predict(self, pack: List[str], deck: List[str]) -> 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() | |
| logits = vals.squeeze(0).cpu().numpy() | |
| return { | |
| "pick": pack[scores.argmax()], | |
| "logits": logits.tolist(), | |
| "scores": scores.tolist(), | |
| } | |
| 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) | |
| vals = self.predict(pack=keys, deck=None)["logits"] | |
| return keys, vals | |