import re import os import pickle import numpy as np import tqdm import torch from torch.utils.data import DataLoader from src.utils import utils from src.training import models # --- Card type / rarity encoding --- MAJOR_TYPES = ['Creature', 'Instant', 'Sorcery', 'Artifact', 'Enchantment', 'Planeswalker', 'Land', 'Battle', 'Tribal'] SUPERTYPES = ['Legendary', 'Basic', 'Snow', 'World'] RARITIES = ['common', 'uncommon', 'rare', 'mythic', 'special'] def type_to_vector(type_line): """Multi-hot over major types (9) and supertypes (4) — 13 dims total.""" v = np.zeros(len(MAJOR_TYPES) + len(SUPERTYPES)) for i, t in enumerate(MAJOR_TYPES): if t in type_line: v[i] = 1.0 for i, s in enumerate(SUPERTYPES): if s in type_line: v[len(MAJOR_TYPES) + i] = 1.0 return v def rarity_to_vector(rarity): """One-hot over rarities (common/uncommon/rare/mythic/special) — 5 dims.""" v = np.zeros(len(RARITIES)) r = rarity.lower().replace(' ', '') for i, name in enumerate(RARITIES): if r.startswith(name): v[i] = 1.0 break return v # --- Feature extraction --- def get_card_features(card): if 'card_faces' in card: name = f"{card['card_faces'][0]['name']} // {card['card_faces'][1]['name']}" colours = str(card['color_identity']) mana_cost = f"{card['card_faces'][0]['mana_cost']} // {card['card_faces'][1]['mana_cost']}" types = f"{card['card_faces'][0]['type_line']} // {card['card_faces'][1]['type_line']}" expansion = str(card['set']) rarity = str(card['rarity']) power = (str(card['card_faces'][0].get('power', 'NaN')) + ' // ' + str(card['card_faces'][1].get('power', 'NaN'))) toughness = (str(card['card_faces'][0].get('toughness', 'NaN')) + ' // ' + str(card['card_faces'][1].get('toughness', 'NaN'))) loyalty = (str(card['card_faces'][0].get('loyalty', 'NaN')) + ' // ' + str(card['card_faces'][1].get('loyalty', 'NaN'))) text = f"{card['card_faces'][0]['oracle_text']} // {card['card_faces'][1]['oracle_text']}" else: name = card['name'] colours = card['color_identity'] mana_cost = card.get('mana_cost', 'NaN') types = card['type_line'] expansion = card['set'] rarity = card['rarity'] power = card.get('power', 'NaN') toughness = card.get('toughness', 'NaN') loyalty = card.get('loyalty', 'NaN') text = card['oracle_text'] return name, colours, mana_cost, types, expansion, rarity, power, toughness, loyalty, text def colour_to_array(colours): v = np.zeros(5) for i, c in enumerate('WUBRG'): if c in colours: v[i] = 1.0 return v def mana_cost_to_array(cost): mana = np.zeros(8) if not cost: return mana for part in cost.split('}'): part = part.replace('{', '') try: mana[0] += int(part) except ValueError: for i, c in enumerate('WUBRGSC'): if c in part: mana[i + 1] += 1 return mana def card_features_to_vector(features): """Numeric feature vector: 5+8+13+5+3 = 34 dims.""" name, colours, mana_cost, types, expansion, rarity, power, toughness, loyalty, text = features colour_v = colour_to_array(colours) # 5 mana_v = mana_cost_to_array(mana_cost.split(' // ')[0]) # 8 type_v = type_to_vector(types.split(' // ')[0]) # 13 rarity_v = rarity_to_vector(rarity) # 5 def _parse_num(val): val = str(val).split(' // ')[0] if val in ('NaN', '*', ''): return 0.0 try: return float(val) except ValueError: nums = re.findall(r'\d+', val) return float(nums[0]) if nums else 0.0 stats = np.array([_parse_num(power), _parse_num(toughness), _parse_num(loyalty)]) # 3 return np.concatenate([colour_v, mana_v, type_v, rarity_v, stats]) def card_to_text(card): name, colours, mana_cost, types, expansion, rarity, power, toughness, loyalty, text = get_card_features(card) return (f'Name: {name} Colours: {colours} Mana Cost: {mana_cost} Types: {types} ' f'Expansion: {expansion} Rarity: {rarity} Power: {power} ' f'Toughness: {toughness} Loyalty: {loyalty} Text: {text}') # --- Embedding helpers --- def normalize_embedding(embedding_path): embedding_dict = utils.get_embedding_dict(embedding_path) tensors = np.array([v for k, v in embedding_dict.items() if k != 'tensor_size']) mean = np.mean(tensors, axis=0) std = np.std(tensors, axis=0) std[std == 0] = 1.0 stats_path = embedding_path.rstrip('.pt') + '_mean_std.pt' with open(stats_path, 'wb') as f: pickle.dump({'mean': mean, 'std': std}, f) for k, v in embedding_dict.items(): if k != 'tensor_size': embedding_dict[k] = (v - mean) / std utils.dump_embedding_dict(embedding_dict, embedding_path.rstrip('.pt') + '_normalized.pt') def fill_embeddings(cards, card_encodings, data): """Add any card names in `cards` missing from `card_encodings` using Scryfall data.""" data_by_name = {c['name']: c for c in data} split_keys = {c.split('//')[0].strip() for c in card_encodings if '//' in c} for card in cards: card = card.replace('_', ' ') if card in card_encodings or card in split_keys: continue try: c = data_by_name[card] except KeyError: card_alt = card.replace("Sol'kanar", "Sol'Kanar") if card_alt.startswith('A-'): card_alt = card_alt[2:] c = next((x for x in data if x['name'] == card_alt), None) if c is None: c = next((x for x in data if x['name'].split('//')[0].strip() == card), None) if c is not None: card_encodings[card] = c def combine_all_embeddings(folder, out_path): all_data = {} for file in os.listdir(folder): if file.endswith('_embedding.pt'): all_data.update(utils.get_embedding_dict(f'{folder}/{file}')) utils.dump_embedding_dict(all_data, out_path) def collate_fn_dict(batch): out = [] for positive, negative, anchor, *_ in batch: if positive: out.append(positive) if isinstance(negative, list): out.extend(negative) elif negative: out.append(negative) if isinstance(anchor, list): out.extend(anchor) elif anchor: out.append(anchor) return out def create_set_embedding_scryfall_language(set_tags, embedding_fn, out_folder=None, card_keys=None, number_vector=False): """Create LLM-based card embeddings, optionally prepending numeric features.""" file_name = ''.join(set_tags) if out_folder is None: out_folder = f'embeddings/scryfall/{file_name}/' os.makedirs(out_folder, exist_ok=True) data = utils.get_card_json() set_tags_lower = [s.lower() for s in set_tags] cards = [c for c in data if c['set'].lower() in set_tags_lower] card_encodings = {card['name']: card for card in cards} if card_keys: fill_embeddings(card_keys, card_encodings, data) names, texts, numbers = [], [], [] for name, card in card_encodings.items(): names.append(name) texts.append(card_to_text(card)) if number_vector: numbers.append(card_features_to_vector(get_card_features(card))) card_encodings = {} if embedding_fn is not None: embedded = embedding_fn(texts).numpy() if number_vector: embedded = np.concatenate([np.array(numbers), embedded], axis=1).astype(np.float64) else: embedded = np.array(numbers).astype(np.float64) for i, name in enumerate(names): card_encodings[name] = embedded[i] print(f'Embedding shape: {embedded.shape}') out_path = os.path.join(out_folder, f'{file_name}_embedding.pt') with open(out_path, 'wb') as f: pickle.dump(card_encodings, f) return card_encodings