maia2 0.1__py2.py3-none-any.whl

maia2/__init__.py

@@ -0,0 +1,3 @@
+"""An amazing sample package for maia2."""
+
+__version__ = "0.1"

maia2/fetch_data.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+
+# Loop over each year
+for year in {2018..2023}; do
+    # Loop over each month
+    for month in {05..12}; do
+        # Form the URL
+        url="https://database.lichess.org/standard/lichess_db_standard_rated_${year}-${month}.pgn.zst"
+
+        # Check if the URL exists
+        if wget --spider $url 2>/dev/null; then
+            # Download the file
+            wget "$url"
+
+            # The downloaded file name
+            filename="lichess_db_standard_rated_${year}-${month}.pgn.zst"
+
+        else
+            echo "File for ${year}-${month} does not exist."
+        fi
+
+        # Stop the loop if we've reached October 2023
+        if [ "$year" -eq 2023 ] && [ "$month" = "10" ]; then
+            break
+        fi
+    done
+
+done

maia2/inference.py

@@ -0,0 +1,154 @@
+from utils import *
+from main import *
+
+
+class TestDataset(torch.utils.data.Dataset):
+    
+    def __init__(self, data, all_moves_dict, elo_dict):
+        
+        self.all_moves_dict = all_moves_dict
+        self.data = data.values.tolist()
+        self.elo_dict = elo_dict
+    
+    def __len__(self):
+        
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        
+        fen, move, elo_self, elo_oppo = self.data[idx]
+
+        if fen.split(' ')[1] == 'w':
+            board = chess.Board(fen)
+        elif fen.split(' ')[1] == 'b':
+            board = chess.Board(fen).mirror()
+            move = mirror_move(move)
+        else:
+            raise ValueError(f"Invalid fen: {fen}")
+            
+        board_input = board_to_tensor(board)
+        
+        elo_self = map_to_category(elo_self, self.elo_dict)
+        elo_oppo = map_to_category(elo_oppo, self.elo_dict)
+        
+        legal_moves, _ = get_side_info(board, move, self.all_moves_dict)
+        
+        return fen, board_input, elo_self, elo_oppo, legal_moves
+
+
+def get_preds(model, dataloader, all_moves_dict_reversed, cfg_inference):
+    
+    all_probs = []
+    predicted_move_probs = []
+    predicted_moves = []
+    predicted_win_probs = []
+    
+    model.eval()
+    with torch.no_grad():
+        
+        for fens, boards, elos_self, elos_oppo, legal_moves in dataloader:
+            
+            if cfg_inference.gpu:
+                boards = boards.cuda()
+                elos_self = elos_self.cuda()
+                elos_oppo = elos_oppo.cuda()
+                legal_moves = legal_moves.cuda()
+
+            logits_maia, _, logits_value = model(boards, elos_self, elos_oppo)
+            logits_maia_legal = logits_maia * legal_moves
+            probs = logits_maia_legal.softmax(dim=-1)
+
+            all_probs.append(probs.cpu())
+            predicted_move_probs.append(probs.max(dim=-1).values.cpu())
+            predicted_move_indices = probs.argmax(dim=-1)
+            for i in range(len(fens)):
+                fen = fens[i]
+                predicted_move = all_moves_dict_reversed[predicted_move_indices[i].item()]
+                if fen.split(' ')[1] == 'b':
+                    predicted_move = mirror_move(predicted_move)
+                predicted_moves.append(predicted_move)
+
+            predicted_win_probs.append((logits_value / 2 + 0.5).cpu())
+    
+    all_probs = torch.cat(all_probs).cpu().numpy()
+    predicted_move_probs = torch.cat(predicted_move_probs).numpy()
+    predicted_win_probs = torch.cat(predicted_win_probs).numpy()
+    
+    return all_probs, predicted_move_probs, predicted_moves, predicted_win_probs
+
+
+def inference_batch(data):
+    
+    cfg = parse_args()
+    cfg_inference = parse_inference_args()
+    if cfg_inference.verbose:
+        show_cfg(cfg_inference)
+
+    all_moves = get_all_possible_moves()
+    all_moves_dict = {move: i for i, move in enumerate(all_moves)}
+    elo_dict = create_elo_dict()
+
+    model = MAIA2Model(len(all_moves), elo_dict, cfg)
+    model = nn.DataParallel(model)
+    
+    checkpoint = torch.load(cfg_inference.model_path, map_location='cpu')
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model = model.module
+    
+    if cfg_inference.gpu:
+        model = model.cuda()
+    
+    all_moves_dict_reversed = {v: k for k, v in all_moves_dict.items()}
+    dataset = TestDataset(data, all_moves_dict, elo_dict)
+    dataloader = torch.utils.data.DataLoader(dataset, 
+                                            batch_size=cfg_inference.batch_size, 
+                                            shuffle=False, 
+                                            drop_last=False,
+                                            num_workers=cfg_inference.num_workers)
+    if cfg_inference.verbose:
+        dataloader = tqdm.tqdm(dataloader)
+    all_probs, predicted_move_probs, predicted_moves, predicted_win_probs = get_preds(model, dataloader, all_moves_dict_reversed, cfg_inference)
+    
+    data['predicted_move'] = predicted_moves
+    data['predicted_move_prob'] = predicted_move_probs
+    data['predicted_win_prob'] = predicted_win_probs
+    data['all_probs'] = all_probs.tolist()
+    
+    return data
+
+def parse_inference_args(args=None):
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--verbose', default=1, type=int)
+    parser.add_argument('--batch_size', default=2048, type=int)
+    parser.add_argument('--num_workers', default=16, type=int)
+    parser.add_argument('--model_path', default='your_model_path', type=str)
+    parser.add_argument('--gpu', default=True, type=bool)
+
+    return parser.parse_args(args)
+
+def show_cfg(cfg):
+    print('Configurations:', flush=True)
+    for arg in vars(cfg):
+        print(f'\t{arg}: {getattr(cfg, arg)}', flush=True)
+
+
+if __name__ == '__main__':
+    
+    data = pd.read_csv('../data/all_reduced_rapid.csv')
+    data = data[data.move_ply > 10]
+    
+    data_novice = data[data['rounded_elo'] <= 1500][['board', 'move', 'active_elo', 'opponent_elo']]
+    data_intermediate = data[(data['rounded_elo'] > 1500) & (data['rounded_elo'] < 2000)][['board', 'move', 'active_elo', 'opponent_elo']]
+    data_advanced = data[data['rounded_elo'] >= 2000][['board', 'move', 'active_elo', 'opponent_elo']]
+    print(f'lens: {len(data_novice)}, {len(data_intermediate)}, {len(data_advanced)}', flush=True)
+    
+    results = []
+    for split in [data_novice, data_intermediate, data_advanced]:
+        inference_batch(split)
+        print(round(len(split[split['predicted_move'] == split['move']]) / len(split), 4))
+    
+    
+    
+    

maia2/main.py

@@ -0,0 +1,742 @@
+import chess.pgn
+import chess
+import pdb
+from multiprocessing import Pool, cpu_count, Queue, Process
+import torch
+import tqdm
+import argparse
+from utils import *
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm.contrib.concurrent import process_map
+import os
+import pandas as pd
+import time
+from einops import rearrange
+
+
+def process_chunks(cfg, pgn_path, pgn_chunks, elo_dict):
+    
+    # process_per_chunk((pgn_chunks[0][0], pgn_chunks[0][1], pgn_path, elo_dict, cfg))
+    
+    if cfg.verbose:
+        results = process_map(process_per_chunk, [(start, end, pgn_path, elo_dict, cfg) for start, end in pgn_chunks], max_workers=len(pgn_chunks), chunksize=1)
+    else:
+        with Pool(processes=len(pgn_chunks)) as pool:
+            results = pool.map(process_per_chunk, [(start, end, pgn_path, elo_dict, cfg) for start, end in pgn_chunks])
+    
+    ret = []
+    count = 0
+    list_of_dicts = []
+    for result, game_count, frequency in results:
+        ret.extend(result)
+        count += game_count
+        list_of_dicts.append(frequency)
+    
+    total_counts = {}
+
+    for d in list_of_dicts:
+        for key, value in d.items():
+            total_counts[key] = total_counts.get(key, 0) + value
+
+    print(total_counts, flush=True)
+    
+    return ret, count, len(pgn_chunks)
+
+
+def process_per_game(game, white_elo, black_elo, white_win, cfg):
+
+    ret = []
+    
+    board = game.board()
+    moves = list(game.mainline_moves())
+    
+    for i, node in enumerate(game.mainline()):
+        
+        move = moves[i]
+        
+        if i >= cfg.first_n_moves:
+            
+            comment = node.comment
+            clock_info = extract_clock_time(comment)
+            
+            if i % 2 == 0:
+                board_input = board.fen()
+                move_input = move.uci()
+                elo_self = white_elo
+                elo_oppo = black_elo
+                active_win = white_win
+
+            else:
+                board_input = board.mirror().fen()
+                move_input = mirror_move(move.uci())
+                elo_self = black_elo
+                elo_oppo = white_elo
+                active_win = - white_win
+
+            if clock_info > cfg.clock_threshold:
+                ret.append((board_input, move_input, elo_self, elo_oppo, active_win))
+        
+        board.push(move)
+        if i == cfg.max_ply:
+            break
+    
+    return ret
+
+
+def game_filter(game):
+    
+    white_elo = game.headers.get("WhiteElo", "?")
+    black_elo = game.headers.get("BlackElo", "?")
+    time_control = game.headers.get("TimeControl", "?")
+    result = game.headers.get("Result", "?")
+    event = game.headers.get("Event", "?")
+    
+    if white_elo == "?" or black_elo == "?" or time_control == "?" or result == "?" or event == "?":
+        return
+
+    if 'Rated' not in event:
+        return
+    
+    if 'Rapid' not in event:
+        return
+    
+    for _, node in enumerate(game.mainline()):
+        if 'clk' not in node.comment:
+            return
+    
+    white_elo = int(white_elo)
+    black_elo = int(black_elo)
+
+    if result == '1-0':
+        white_win = 1
+    elif result == '0-1':
+        white_win = -1
+    elif result == '1/2-1/2':
+        white_win = 0
+    else:
+        return
+    
+    return game, white_elo, black_elo, white_win
+
+
+def process_per_chunk(args):
+
+    start_pos, end_pos, pgn_path, elo_dict, cfg = args
+    
+    ret = []
+    game_count = 0
+    
+    frequency = {}
+    
+    with open(pgn_path, 'r', encoding='utf-8') as pgn_file:
+        
+        pgn_file.seek(start_pos)
+
+        while pgn_file.tell() < end_pos:
+            
+            game = chess.pgn.read_game(pgn_file)
+            
+            if game is None:
+                break
+
+            filtered_game = game_filter(game)
+            if filtered_game:
+                game, white_elo, black_elo, white_win = filtered_game
+                white_elo = map_to_category(white_elo, elo_dict)
+                black_elo = map_to_category(black_elo, elo_dict)
+                
+                if white_elo < black_elo:
+                    range_1, range_2 = black_elo, white_elo
+                else:
+                    range_1, range_2 = white_elo, black_elo
+                
+                freq = frequency.get((range_1, range_2), 0)
+                if freq >= cfg.max_games_per_elo_range:
+                    continue
+                
+                ret_per_game = process_per_game(game, white_elo, black_elo, white_win, cfg)
+                ret.extend(ret_per_game)
+                if len(ret_per_game):
+                    
+                    if (range_1, range_2) in frequency:
+                        frequency[(range_1, range_2)] += 1
+                    else:
+                        frequency[(range_1, range_2)] = 1
+                    
+                    game_count += 1
+    
+    return ret, game_count, frequency
+
+
+class MAIA1Dataset(torch.utils.data.Dataset):
+    
+    def __init__(self, data, all_moves_dict, elo_dict, cfg):
+        
+        self.all_moves_dict = all_moves_dict
+        self.cfg = cfg
+        self.data = data.values.tolist()
+        self.elo_dict = elo_dict
+    
+    def __len__(self):
+        
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        
+        fen, move, elo_self, elo_oppo, white_active = self.data[idx]
+
+        if white_active:
+            board = chess.Board(fen)
+        else:
+            board = chess.Board(fen).mirror()
+            move = mirror_move(move)
+            
+        board_input = board_to_tensor(board)
+        move_input = self.all_moves_dict[move]
+        
+        elo_self = map_to_category(elo_self, self.elo_dict)
+        elo_oppo = map_to_category(elo_oppo, self.elo_dict)
+        
+        legal_moves, side_info = get_side_info(board, move, self.all_moves_dict)
+        
+        return board_input, move_input, elo_self, elo_oppo, legal_moves, side_info
+
+
+class MAIA2Dataset(torch.utils.data.Dataset):
+    
+    
+    def __init__(self, data, all_moves_dict, cfg):
+        
+        self.all_moves_dict = all_moves_dict
+        self.data = data
+        self.cfg = cfg
+    
+    def __len__(self):
+        
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        
+        board_input, move_uci, elo_self, elo_oppo, active_win = self.data[idx]
+        
+        board = chess.Board(board_input)
+        board_input = board_to_tensor(board)
+        
+        legal_moves, side_info = get_side_info(board, move_uci, self.all_moves_dict)
+        
+        move_input = self.all_moves_dict[move_uci]
+        
+        return board_input, move_input, elo_self, elo_oppo, legal_moves, side_info, active_win
+
+
+class BasicBlock(torch.nn.Module):
+
+    def __init__(self, in_planes, planes, stride=1):
+        super(BasicBlock, self).__init__()
+        
+        mid_planes = planes
+        
+        self.conv1 = torch.nn.Conv2d(in_planes, mid_planes, kernel_size=3, stride=stride, padding=1, bias=False)
+        self.bn1 = torch.nn.BatchNorm2d(mid_planes)
+        self.conv2 = torch.nn.Conv2d(mid_planes, planes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn2 = torch.nn.BatchNorm2d(planes)
+        self.dropout = nn.Dropout(p=0.5)
+
+    def forward(self, x):
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = F.relu(out)
+        out = self.dropout(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out += x
+        out = F.relu(out)
+
+        return out
+
+
+class ChessResNet(torch.nn.Module):
+    
+    def __init__(self, block, cfg):
+        super(ChessResNet, self).__init__()
+        
+        self.conv1 = torch.nn.Conv2d(cfg.input_channels, cfg.dim_cnn, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn1 = torch.nn.BatchNorm2d(cfg.dim_cnn)
+        self.layers = self._make_layer(block, cfg.dim_cnn, cfg.num_blocks_cnn)
+        self.conv_last = torch.nn.Conv2d(cfg.dim_cnn, cfg.vit_length, kernel_size=3, stride=1, padding=1, bias=False)
+        self.bn_last = torch.nn.BatchNorm2d(cfg.vit_length)
+
+    def _make_layer(self, block, planes, num_blocks, stride=1):
+        
+        layers = []
+        for _ in range(num_blocks):
+            layers.append(block(planes, planes, stride))
+        
+        return torch.nn.Sequential(*layers)
+
+    def forward(self, x):
+        
+        out = F.relu(self.bn1(self.conv1(x)))
+        out = self.layers(out)
+        out = self.conv_last(out)
+        out = self.bn_last(out)
+        
+        return out
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim, dropout = 0.):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, dim),
+            nn.Dropout(dropout)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = nn.Softmax(dim = -1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+
+class EloAwareAttention(nn.Module):
+    def __init__(self, dim, heads=8, dim_head=64, dropout=0., elo_dim=64):
+        super().__init__()
+        inner_dim = dim_head * heads
+        project_out = not (heads == 1 and dim_head == dim)
+
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = nn.Softmax(dim=-1)
+        self.dropout = nn.Dropout(dropout)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
+
+        self.elo_query = nn.Linear(elo_dim, inner_dim, bias=False)
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, dim),
+            nn.Dropout(dropout)
+        ) if project_out else nn.Identity()
+
+    def forward(self, x, elo_emb):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim=-1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)
+
+        elo_effect = self.elo_query(elo_emb).view(x.size(0), self.heads, 1, -1)
+        q = q + elo_effect
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+        attn = self.dropout(attn)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., elo_dim=64):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = nn.ModuleList([])
+        self.elo_layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.elo_layers.append(nn.ModuleList([
+                EloAwareAttention(dim, heads = heads, dim_head = dim_head, dropout = dropout, elo_dim = elo_dim),
+                FeedForward(dim, mlp_dim, dropout = dropout)
+            ]))
+
+    def forward(self, x, elo_emb):
+        for attn, ff in self.elo_layers:
+            x = attn(x, elo_emb) + x
+            x = ff(x) + x
+
+        return self.norm(x)
+
+
+class MAIA2Model(torch.nn.Module):
+    
+    def __init__(self, output_dim, elo_dict, cfg):
+        super(MAIA2Model, self).__init__()
+        
+        self.cfg = cfg
+        self.chess_cnn = ChessResNet(BasicBlock, cfg)
+        
+        heads = 16
+        dim_head = 64
+        self.to_patch_embedding = nn.Sequential(
+            nn.Linear(8 * 8, cfg.dim_vit),
+            nn.LayerNorm(cfg.dim_vit),
+        )
+        self.transformer = Transformer(cfg.dim_vit, cfg.num_blocks_vit, heads, dim_head, mlp_dim=cfg.dim_vit, dropout = 0.1, elo_dim = cfg.elo_dim * 2)
+        self.pos_embedding = nn.Parameter(torch.randn(1, cfg.vit_length, cfg.dim_vit))
+        
+        self.fc_1 = nn.Linear(cfg.dim_vit, output_dim)
+        # self.fc_1_1 = nn.Linear(cfg.dim_vit, cfg.dim_vit)
+        self.fc_2 = nn.Linear(cfg.dim_vit, output_dim + 6 + 6 + 1 + 64 + 64)
+        # self.fc_2_1 = nn.Linear(cfg.dim_vit, cfg.dim_vit)
+        self.fc_3 = nn.Linear(128, 1)
+        self.fc_3_1 = nn.Linear(cfg.dim_vit, 128)
+        
+        self.elo_embedding = torch.nn.Embedding(len(elo_dict), cfg.elo_dim)
+        
+        self.dropout = nn.Dropout(p=0.1)
+        self.last_ln = nn.LayerNorm(cfg.dim_vit)
+
+
+    def forward(self, boards, elos_self, elos_oppo):
+        
+        batch_size = boards.size(0)
+        boards = boards.view(batch_size, self.cfg.input_channels, 8, 8)
+        embs = self.chess_cnn(boards)
+        embs = embs.view(batch_size, embs.size(1), 8 * 8)
+        x = self.to_patch_embedding(embs)
+        x += self.pos_embedding
+        x = self.dropout(x)
+        
+        elos_emb_self = self.elo_embedding(elos_self)
+        elos_emb_oppo = self.elo_embedding(elos_oppo)
+        elos_emb = torch.cat((elos_emb_self, elos_emb_oppo), dim=1)
+        x = self.transformer(x, elos_emb).mean(dim=1)
+        
+        x = self.last_ln(x)
+
+        logits_maia = self.fc_1(x)
+        logits_side_info = self.fc_2(x)
+        logits_value = self.fc_3(self.dropout(torch.relu(self.fc_3_1(x)))).squeeze(dim=-1)
+        
+        return logits_maia, logits_side_info, logits_value
+
+
+def read_monthly_data_path(cfg):
+    
+    print('Training Data:', flush=True)
+    pgn_paths = []
+    
+    for year in range(cfg.start_year, cfg.end_year + 1):
+        start_month = cfg.start_month if year == cfg.start_year else 1
+        end_month = cfg.end_month if year == cfg.end_year else 12
+
+        for month in range(start_month, end_month + 1):
+            formatted_month = f"{month:02d}"
+            pgn_path = cfg.data_root + f"/lichess_db_standard_rated_{year}-{formatted_month}.pgn"
+            # skip 2019-12
+            if year == 2019 and month == 12:
+                continue
+            print(pgn_path, flush=True)
+            pgn_paths.append(pgn_path)
+            
+    return pgn_paths
+
+
+def evaluate(model, dataloader):
+    
+    counter = 0
+    correct_move = 0
+    
+    model.eval()
+    with torch.no_grad():
+        
+        for boards, labels, elos_self, elos_oppo, legal_moves, side_info in dataloader:
+            
+            boards = boards.cuda()
+            labels = labels.cuda()
+            elos_self = elos_self.cuda()
+            elos_oppo = elos_oppo.cuda()
+            legal_moves = legal_moves.cuda()
+
+            logits_maia, logits_side_info, logits_value = model(boards, elos_self, elos_oppo)
+            logits_maia_legal = logits_maia * legal_moves
+            preds = logits_maia_legal.argmax(dim=-1)
+            correct_move += (preds == labels).sum().item()
+            
+            counter += len(labels)
+
+    return correct_move, counter
+
+
+def evaluate_MAIA1_data(model, all_moves_dict, elo_dict, cfg, tiny=False):
+    
+    elo_list = range(1000, 2600, 100)
+
+    for i in elo_list:
+        start = i
+        end = i + 100
+        file_path = f"../data/test/KDDTest_{start}-{end}.csv"
+        data = pd.read_csv(file_path)
+        data = data[data.type == 'Rapid'][['board', 'move', 'active_elo', 'opponent_elo', 'white_active']]
+        dataset = MAIA1Dataset(data, all_moves_dict, elo_dict, cfg)
+        dataloader = torch.utils.data.DataLoader(dataset, 
+                                                batch_size=cfg.batch_size, 
+                                                shuffle=False, 
+                                                drop_last=False,
+                                                num_workers=cfg.num_workers)
+        if cfg.verbose:
+            dataloader = tqdm.tqdm(dataloader)
+        print(f'Testing Elo Range {start}-{end} with MAIA 1 data:', flush=True)
+        correct_move, counter = evaluate(model, dataloader)
+        print(f'Accuracy Move Prediction: {round(correct_move / counter, 4)}', flush=True)
+        if tiny:
+            break
+
+
+def train_chunks(cfg, data, model, optimizer, all_moves_dict, criterion_maia, criterion_side_info, criterion_value):
+    
+    dataset_train = MAIA2Dataset(data, all_moves_dict, cfg)
+    dataloader_train = torch.utils.data.DataLoader(dataset_train, 
+                                                    batch_size=cfg.batch_size, 
+                                                    shuffle=True, 
+                                                    drop_last=False,
+                                                    num_workers=cfg.num_workers)
+    if cfg.verbose:
+        dataloader_train = tqdm.tqdm(dataloader_train)
+    
+    avg_loss = 0
+    avg_loss_maia = 0
+    avg_loss_side_info = 0
+    avg_loss_value = 0
+    step = 0
+    for boards, labels, elos_self, elos_oppo, legal_moves, side_info, wdl in dataloader_train:
+
+        model.train()
+        boards = boards.cuda()
+        labels = labels.cuda()
+        elos_self = elos_self.cuda()
+        elos_oppo = elos_oppo.cuda()
+        side_info = side_info.cuda()
+        wdl = wdl.float().cuda()
+        
+        logits_maia, logits_side_info, logits_value = model(boards, elos_self, elos_oppo)
+        
+        loss = 0
+        loss_maia = criterion_maia(logits_maia, labels)
+        loss += loss_maia
+        
+        if cfg.side_info:
+        
+            loss_side_info = criterion_side_info(logits_side_info, side_info) * cfg.side_info_coefficient
+            loss += loss_side_info
+        
+        if cfg.value:
+            loss_value = criterion_value(logits_value, wdl) * cfg.value_coefficient
+            loss += loss_value
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        
+        avg_loss += loss.item()
+        avg_loss_maia += loss_maia.item()
+        if cfg.side_info:
+            avg_loss_side_info += loss_side_info.item()
+        if cfg.value:
+            avg_loss_value += loss_value.item()
+        step += 1
+    
+    return round(avg_loss / step, 3), round(avg_loss_maia / step, 3), round(avg_loss_side_info / step, 3), round(avg_loss_value / step, 3)
+
+
+def preprocess_thread(queue, cfg, pgn_path, pgn_chunks_sublist, elo_dict):
+    
+    data, game_count, chunk_count = process_chunks(cfg, pgn_path, pgn_chunks_sublist, elo_dict)
+    queue.put([data, game_count, chunk_count])
+    del data
+
+
+def worker_wrapper(semaphore, *args, **kwargs):
+    with semaphore:
+        preprocess_thread(*args, **kwargs)
+
+
+def parse_args(args=None):
+
+    parser = argparse.ArgumentParser()
+
+    # Supporting Arguments
+    parser.add_argument('--data_root', default='your_data_root', type=str)
+    parser.add_argument('--seed', default=42, type=int)
+    parser.add_argument('--num_workers', default=16, type=int)
+    parser.add_argument('--verbose', default=1, type=int)
+    parser.add_argument('--max_epochs', default=3, type=int)
+    parser.add_argument('--max_ply', default=300, type=int)
+    parser.add_argument('--clock_threshold', default=30, type=int)
+    parser.add_argument('--chunk_size', default=20000, type=str)
+    parser.add_argument('--start_year', default=2018, type=int)
+    parser.add_argument('--start_month', default=5, type=int)
+    parser.add_argument('--end_year', default=2023, type=int)
+    parser.add_argument('--end_month', default=11, type=int)
+    parser.add_argument('--from_checkpoint', default=False, type=bool)
+    parser.add_argument('--checkpoint_epoch', default=2, type=int)
+    parser.add_argument('--checkpoint_year', default=2018, type=int)
+    parser.add_argument('--checkpoint_month', default=6, type=int)
+    parser.add_argument('--num_cpu_left', default=16, type=int)
+    parser.add_argument('--queue_length', default=2, type=int)
+    parser.add_argument('--max_games_per_elo_range', default=20, type=int)
+
+    # Tunable Arguments
+    parser.add_argument('--lr', default=1e-4, type=float)
+    parser.add_argument('--wd', default=1e-5, type=float)
+    parser.add_argument('--batch_size', default=8192, type=int)
+    parser.add_argument('--first_n_moves', default=10, type=int)
+    parser.add_argument('--last_n_moves', default=10, type=int)
+    parser.add_argument('--dim_cnn', default=256, type=int)
+    parser.add_argument('--dim_vit', default=1024, type=int)
+    parser.add_argument('--num_blocks_cnn', default=5, type=int)
+    parser.add_argument('--num_blocks_vit', default=2, type=int)
+    parser.add_argument('--input_channels', default=18, type=int)
+    parser.add_argument('--vit_length', default=8, type=int)
+    parser.add_argument('--elo_dim', default=128, type=int)
+    parser.add_argument('--side_info', default=True, type=bool)
+    parser.add_argument('--side_info_coefficient', default=1, type=float)
+    parser.add_argument('--value', default=True, type=bool)
+    parser.add_argument('--value_coefficient', default=1, type=float)
+
+    return parser.parse_args(args)
+
+
+if __name__ == '__main__':
+    
+    cfg = parse_args()
+    print('Configurations:', flush=True)
+    for arg in vars(cfg):
+        print(f'\t{arg}: {getattr(cfg, arg)}', flush=True)
+    seed_everything(cfg.seed)
+    num_processes = cpu_count() - cfg.num_cpu_left
+    
+    save_root = f'../tmp/{cfg.lr}_{cfg.batch_size}_{cfg.wd}_MAIA2/'
+    if not os.path.exists(save_root):
+        os.makedirs(save_root)
+
+    all_moves = get_all_possible_moves()
+    all_moves_dict = {move: i for i, move in enumerate(all_moves)}
+    elo_dict = create_elo_dict()
+
+    model = MAIA2Model(len(all_moves), elo_dict, cfg)
+
+    print(model, flush=True)
+    model = model.cuda()
+    model = nn.DataParallel(model)
+    criterion_maia = nn.CrossEntropyLoss()
+    criterion_side_info = nn.BCEWithLogitsLoss()
+    criterion_value = nn.MSELoss()
+    
+    optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.lr, weight_decay=cfg.wd)
+    N_params = count_parameters(model)
+    print(f'Trainable Parameters: {N_params}', flush=True)
+    
+    accumulated_samples = 0
+    accumulated_games = 0
+    
+    if cfg.from_checkpoint:
+        formatted_month = f"{cfg.checkpoint_month:02d}"
+        checkpoint = torch.load(save_root + f'epoch_{cfg.checkpoint_epoch}_{cfg.checkpoint_year}-{formatted_month}.pgn.pt')
+        model.load_state_dict(checkpoint['model_state_dict'])
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        accumulated_samples = checkpoint['accumulated_samples']
+        accumulated_games = checkpoint['accumulated_games']
+
+    for epoch in range(cfg.max_epochs):
+        
+        print(f'Epoch {epoch + 1}', flush=True)
+        pgn_paths = read_monthly_data_path(cfg)
+        
+        num_file = 0
+        for pgn_path in pgn_paths:
+            
+            start_time = time.time()
+            decompress_zst(pgn_path + '.zst', pgn_path)
+            print(f'Decompressing {pgn_path} took {readable_time(time.time() - start_time)}', flush=True)
+
+            pgn_chunks = read_or_create_chunks(pgn_path, cfg)
+            print(f'Training {pgn_path} with {len(pgn_chunks)} chunks.', flush=True)
+            
+            queue = Queue(maxsize=cfg.queue_length)
+            
+            pgn_chunks_sublists = []
+            for i in range(0, len(pgn_chunks), num_processes):
+                pgn_chunks_sublists.append(pgn_chunks[i:i + num_processes])
+            
+            pgn_chunks_sublist = pgn_chunks_sublists[0]
+            # For debugging only
+            # process_chunks(cfg, pgn_path, pgn_chunks_sublist, elo_dict)
+            worker = Process(target=preprocess_thread, args=(queue, cfg, pgn_path, pgn_chunks_sublist, elo_dict))
+            worker.start()
+            
+            num_chunk = 0
+            offset = 0
+            while True:
+                if not queue.empty():
+                    if offset + 1 < len(pgn_chunks_sublists):
+                        pgn_chunks_sublist = pgn_chunks_sublists[offset + 1]
+                        worker = Process(target=preprocess_thread, args=(queue, cfg, pgn_path, pgn_chunks_sublist, elo_dict))
+                        worker.start()
+                        offset += 1
+                    data, game_count, chunk_count = queue.get()
+                    loss, loss_maia, loss_side_info, loss_value = train_chunks(cfg, data, model, optimizer, all_moves_dict, criterion_maia, criterion_side_info, criterion_value)
+                    num_chunk += chunk_count
+                    accumulated_samples += len(data)
+                    accumulated_games += game_count
+                    print(f'[{num_chunk}/{len(pgn_chunks)}]', flush=True)
+                    print(f'[# Positions]: {readable_num(accumulated_samples)}', flush=True)
+                    print(f'[# Games]: {readable_num(accumulated_games)}', flush=True)
+                    print(f'[# Loss]: {loss} | [# Loss MAIA]: {loss_maia} | [# Loss Side Info]: {loss_side_info} | [# Loss Value]: {loss_value}', flush=True)
+                    if num_chunk == len(pgn_chunks):
+                        break
+
+            num_file += 1
+            print(f'### ({num_file} / {len(pgn_paths)}) Took {readable_time(time.time() - start_time)} to train {pgn_path} with {len(pgn_chunks)} chunks.', flush=True)
+            os.remove(pgn_path)
+            
+            # evaluate_MAIA1_data(model, all_moves_dict, elo_dict, cfg, tiny=False)
+            torch.save({'model_state_dict': model.state_dict(),
+                        'optimizer_state_dict': optimizer.state_dict(),
+                        'accumulated_samples': accumulated_samples,
+                        'accumulated_games': accumulated_games}, f'{save_root}epoch_{epoch + 1}_{pgn_path[-11:]}.pt')

maia2/utils.py

@@ -0,0 +1,347 @@
+import pdb
+import chess
+import pickle
+import os
+import random
+import numpy as np
+import torch
+import time
+import requests
+import tqdm
+import pyzstd
+import re
+
+
+def seed_everything(seed: int):
+
+    random.seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def delete_file(filename):
+    
+    if os.path.exists(filename):
+        os.remove(filename)
+        print(f"Data {filename} has been deleted.")
+    else:
+        print(f"The file '{filename}' does not exist.")
+
+
+def readable_num(num):
+    
+    if num >= 1e9:  # if parameters are in the billions
+        return f'{num / 1e9:.2f}B'
+    elif num >= 1e6:  # if parameters are in the millions
+        return f'{num / 1e6:.2f}M'
+    elif num >= 1e3:  # if parameters are in the thousands
+        return f'{num / 1e3:.2f}K'
+    else:
+        return str(num)
+
+
+def readable_time(elapsed_time):
+
+    hours, rem = divmod(elapsed_time, 3600)
+    minutes, seconds = divmod(rem, 60)
+
+    if hours > 0:
+        return f"{int(hours)}h {int(minutes)}m {seconds:.2f}s"
+    elif minutes > 0:
+        return f"{int(minutes)}m {seconds:.2f}s"
+    else:
+        return f"{seconds:.2f}s"
+
+
+def count_parameters(model):
+    
+    total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    
+    return readable_num(total_params)
+
+
+def create_elo_dict():
+    
+    inteval = 100
+    start = 1100
+    end = 2000
+    
+    range_dict = {f"<{start}": 0}
+    range_index = 1
+
+    for lower_bound in range(start, end - 1, inteval):
+        upper_bound = lower_bound + inteval
+        range_dict[f"{lower_bound}-{upper_bound - 1}"] = range_index
+        range_index += 1
+
+    range_dict[f">={end}"] = range_index
+    
+    # print(range_dict, flush=True)
+    
+    return range_dict
+
+
+def map_to_category(elo, elo_dict):
+
+    inteval = 100
+    start = 1100
+    end = 2000
+    
+    if elo < start:
+        return elo_dict[f"<{start}"]
+    elif elo >= end:
+        return elo_dict[f">={end}"]
+    else:
+        for lower_bound in range(start, end - 1, inteval):
+            upper_bound = lower_bound + inteval
+            if lower_bound <= elo < upper_bound:
+                return elo_dict[f"{lower_bound}-{upper_bound - 1}"]
+
+
+def get_side_info(board, move_uci, all_moves_dict):
+    move = chess.Move.from_uci(move_uci)
+    
+    moving_piece = board.piece_at(move.from_square)
+    captured_piece = board.piece_at(move.to_square)
+
+    from_square_encoded = torch.zeros(64)
+    from_square_encoded[move.from_square] = 1
+
+    to_square_encoded = torch.zeros(64)
+    to_square_encoded[move.to_square] = 1
+    
+    if move_uci == 'e1g1':
+        rook_move = chess.Move.from_uci('h1f1')
+        from_square_encoded[rook_move.from_square] = 1
+        to_square_encoded[rook_move.to_square] = 1
+    
+    if move_uci == 'e1c1':
+        rook_move = chess.Move.from_uci('a1d1')
+        from_square_encoded[rook_move.from_square] = 1
+        to_square_encoded[rook_move.to_square] = 1
+
+    board.push(move)
+    is_check = board.is_check()
+    board.pop()
+    
+    # Order: Pawn, Knight, Bishop, Rook, Queen, King
+    side_info = torch.zeros(6 + 6 + 1)
+    side_info[moving_piece.piece_type - 1] = 1
+    if move_uci in ['e1g1', 'e1c1']:
+        side_info[3] = 1
+    if captured_piece:
+        side_info[6 + captured_piece.piece_type - 1] = 1
+    if is_check:
+        side_info[-1] = 1
+    
+    legal_moves = torch.zeros(len(all_moves_dict))
+    legal_moves_idx = torch.tensor([all_moves_dict[move.uci()] for move in board.legal_moves])
+    legal_moves[legal_moves_idx] = 1
+    
+    side_info = torch.cat([side_info, from_square_encoded, to_square_encoded, legal_moves], dim=0)
+    
+    return legal_moves, side_info
+
+
+def extract_clock_time(comment):
+    
+    match = re.search(r'\[%clk (\d+):(\d+):(\d+)\]', comment)
+    if match:
+        hours, minutes, seconds = map(int, match.groups())
+        return hours * 3600 + minutes * 60 + seconds
+    return None
+    
+
+def read_or_create_chunks(pgn_path, cfg):
+
+    cache_file = pgn_path.replace('.pgn', '_chunks.pkl')
+
+    if os.path.exists(cache_file):
+        print(f"Loading cached chunks from {cache_file}")
+        with open(cache_file, 'rb') as f:
+            pgn_chunks = pickle.load(f)
+    else:
+        print(f"Cache not found. Creating chunks for {pgn_path}")
+        start_time = time.time()
+        pgn_chunks = get_chunks(pgn_path, cfg.chunk_size)
+        print(f'Chunking took {readable_time(time.time() - start_time)}', flush=True)
+        
+        with open(cache_file, 'wb') as f:
+            pickle.dump(pgn_chunks, f)
+    
+    return pgn_chunks
+
+
+def board_to_tensor(board):
+    
+    piece_types = [chess.PAWN, chess.KNIGHT, chess.BISHOP, chess.ROOK, chess.QUEEN, chess.KING]
+    num_piece_channels = 12  # 6 piece types * 2 colors
+    additional_channels = 6  # 1 for player's turn, 4 for castling rights, 1 for en passant
+    tensor = torch.zeros((num_piece_channels + additional_channels, 8, 8), dtype=torch.float32)
+
+    # Precompute indices for each piece type
+    piece_indices = {piece: i for i, piece in enumerate(piece_types)}
+
+    # Fill tensor for each piece type
+    for piece_type in piece_types:
+        for color in [True, False]:  # True is White, False is Black
+            piece_map = board.pieces(piece_type, color)
+            index = piece_indices[piece_type] + (0 if color else 6)
+            for square in piece_map:
+                row, col = divmod(square, 8)
+                tensor[index, row, col] = 1.0
+
+    # Player's turn channel (White = 1, Black = 0)
+    turn_channel = num_piece_channels
+    if board.turn == chess.WHITE:
+        tensor[turn_channel, :, :] = 1.0
+
+    # Castling rights channels
+    castling_rights = [board.has_kingside_castling_rights(chess.WHITE),
+                       board.has_queenside_castling_rights(chess.WHITE),
+                       board.has_kingside_castling_rights(chess.BLACK),
+                       board.has_queenside_castling_rights(chess.BLACK)]
+    for i, has_right in enumerate(castling_rights):
+        if has_right:
+            tensor[num_piece_channels + 1 + i, :, :] = 1.0
+
+    # En passant target channel
+    ep_channel = num_piece_channels + 5
+    if board.ep_square is not None:
+        row, col = divmod(board.ep_square, 8)
+        tensor[ep_channel, row, col] = 1.0
+
+    return tensor
+
+def generate_pawn_promotions():
+    # Define the promotion rows for both colors and the promotion pieces
+    # promotion_rows = {'white': '7', 'black': '2'}
+    promotion_rows = {'white': '7'}
+    promotion_pieces = ['q', 'r', 'b', 'n']
+    promotions = []
+
+    # Iterate over each color
+    for color, row in promotion_rows.items():
+        # Target rows for promotion (8 for white, 1 for black)
+        target_row = '8' if color == 'white' else '1'
+
+        # Each file from 'a' to 'h'
+        for file in 'abcdefgh':
+            # Direct move to promotion
+            for piece in promotion_pieces:
+                promotions.append(f'{file}{row}{file}{target_row}{piece}')
+
+            # Capturing moves to the left and right (if not on the edges of the board)
+            if file != 'a':
+                left_file = chr(ord(file) - 1)  # File to the left
+                for piece in promotion_pieces:
+                    promotions.append(f'{file}{row}{left_file}{target_row}{piece}')
+
+            if file != 'h':
+                right_file = chr(ord(file) + 1)  # File to the right
+                for piece in promotion_pieces:
+                    promotions.append(f'{file}{row}{right_file}{target_row}{piece}')
+
+    return promotions
+
+
+def mirror_square(square):
+    
+    file = square[0]
+    rank = str(9 - int(square[1]))
+    
+    return file + rank
+
+
+def mirror_move(move_uci):
+    # Check if the move is a promotion (length of UCI string will be more than 4)
+    is_promotion = len(move_uci) > 4
+
+    # Extract the start and end squares, and the promotion piece if applicable
+    start_square = move_uci[:2]
+    end_square = move_uci[2:4]
+    promotion_piece = move_uci[4:] if is_promotion else ""
+
+    # Mirror the start and end squares
+    mirrored_start = mirror_square(start_square)
+    mirrored_end = mirror_square(end_square)
+
+    # Return the mirrored move, including the promotion piece if applicable
+    return mirrored_start + mirrored_end + promotion_piece
+
+
+def get_chunks(pgn_path, chunk_size):
+
+    chunks = []
+    with open(pgn_path, 'r', encoding='utf-8') as pgn_file:
+        while True:
+            start_pos = pgn_file.tell()
+            game_count = 0
+            while game_count < chunk_size:
+                line = pgn_file.readline()
+                if not line:
+                    break
+                if line[-4:] == "1-0\n" or line[-4:] == "0-1\n":
+                    game_count += 1
+                if line[-8:] == "1/2-1/2\n":
+                    game_count += 1
+                if line[-2:] == "*\n":
+                    game_count += 1
+            line = pgn_file.readline()
+            if line not in ["\n", ""]:
+                raise ValueError
+            end_pos = pgn_file.tell()
+            chunks.append((start_pos, end_pos))
+            if not line:
+                break
+
+    return chunks
+
+
+def decompress_zst(file_path, decompressed_path):
+    """ Decompress a .zst file using pyzstd """
+    with open(file_path, 'rb') as compressed_file, open(decompressed_path, 'wb') as decompressed_file:
+        pyzstd.decompress_stream(compressed_file, decompressed_file)
+
+
+def get_all_possible_moves():
+    
+    all_moves = []
+
+    for rank in range(8):
+        for file in range(8): 
+            square = chess.square(file, rank)
+            
+            board = chess.Board(None)
+            board.set_piece_at(square, chess.Piece(chess.QUEEN, chess.WHITE))
+            legal_moves = list(board.legal_moves)
+            all_moves.extend(legal_moves)
+            
+            board = chess.Board(None)
+            board.set_piece_at(square, chess.Piece(chess.KNIGHT, chess.WHITE))
+            legal_moves = list(board.legal_moves)
+            all_moves.extend(legal_moves)
+    
+    all_moves = [all_moves[i].uci() for i in range(len(all_moves))]
+    
+    pawn_promotions = generate_pawn_promotions()
+    
+    return all_moves + pawn_promotions
+
+
+def chunks(lst, n):
+    """Yield successive n-sized chunks from lst."""
+    for i in range(0, len(lst), n):
+        yield lst[i:i + n]
+
+
+if __name__ == '__main__':
+    
+    pass
+
+

maia2-0.1.dist-info/LICENSE

@@ -0,0 +1,24 @@
+BSD 2-Clause License
+
+Copyright (c) 2024, CSSLab
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

maia2-0.1.dist-info/METADATA

@@ -0,0 +1,10 @@
+Metadata-Version: 2.1
+Name: maia2
+Version: 0.1
+Summary: An amazing sample package for maia2.
+Author-email: csslab <josephtang@cs.toronto.edu>
+Description-Content-Type: text/markdown
+Classifier: License :: OSI Approved :: MIT License
+Project-URL: Home, https://github.com/CSSLab/maia2
+
+# maia2

maia2-0.1.dist-info/RECORD

@@ -0,0 +1,9 @@
+maia2/__init__.py,sha256=8z7obJwpM5l5hGYw8MvEwuViELcuO-D991YkvWhMJh0,64
+maia2/fetch_data.sh,sha256=LefV6tal7pp0a34ODhcx6y1-i0p4m1mLvS1ALIB0PeM,746
+maia2/inference.py,sha256=43rHIrV2oYVLvxYrVF9L_5C4Zad7BjsKGKsVcKPRBuA,5542
+maia2/main.py,sha256=A1YM01iJN2XzGfK1a8vWTB7HbwrjFhgRo2A65qXR6VY,26106
+maia2/utils.py,sha256=dvriII8I6zzS-ZIsNEDDV6tAXfiqaa0OZ0NTOARlgd0,10560
+maia2-0.1.dist-info/LICENSE,sha256=T0nFxRZvtWtfqPnjU6vIWkdL0t13hvAAhwHRZhluelc,1295
+maia2-0.1.dist-info/WHEEL,sha256=Sgu64hAMa6g5FdzHxXv9Xdse9yxpGGMeagVtPMWpJQY,99
+maia2-0.1.dist-info/METADATA,sha256=QWXxHkV-w_G_sm_NgNuFtkyMXsr9hHBm1Yj10JGGf5w,293
+maia2-0.1.dist-info/RECORD,,

maia2-0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: flit 3.9.0
+Root-Is-Purelib: true
+Tag: py2-none-any
+Tag: py3-none-any