Spaces:
Sleeping
Sleeping
| 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 | |