update train script

train.py

@@ -1,9 +1,15 @@
+import re
 import argparse
+import json
 import torch
 import torch.nn as nn
+from tqdm import tqdm
 from torch.nn import functional as F
 from gpt_p.model import DecoderTransformer
+from torch.optim.lr_scheduler import _LRScheduler
+import math
 from datasets import load_dataset
+import wandb
 
 
 torch.manual_seed(420) # 1337
@@ -12,7 +18,7 @@ base_name = 'gpt-p_CHARS_CHAT_'
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 context_size = 256 # how many tokens to consider while generating the next
 batch_size = 128 # how many independent sequences will we process in parallel
-max_iters = 30_000
+max_iters = 50_000
 learning_rate = 3e-5
 eval_interval = 100
 eval_iters = 20 # number evaluation iterations
@@ -21,28 +27,304 @@ n_layer = 6 # number of transformer layers
 n_head = 6
 dropout = 0.2 # dropout factor
 
-dataset = load_dataset('Lichess/standard-chess-games', split='train')
-content = '\n'.join(list(filter(lambda x: 'eval' not in x, dataset['movetext'])))
+mask_all_data = True
+use_scheduler = False
+
+dataset = load_dataset('Lichess/standard-chess-games', '2014-08', split='train')
+og_samples = list(filter(lambda x: 'eval' not in x, dataset['movetext']))
+
+
+new_dataset = load_dataset('Lichess/standard-chess-games', '2024-07', split='train', data_files=[f'data/year=2024/month=07/train-{str(i).zfill(5)}-of-00384.parquet' for i in range(10)])
+
+new_dataset = [re.sub('[0-9]+\.\.\.', '', re.sub('{[^\}]*}', '', foo)).replace('  ', ' ').replace('  ', ' ') for foo in dataset['movetext']]
+
+og_samples += new_dataset
+
+if mask_all_data:
+    content = '\n'.join(list(filter(lambda x: 'eval' not in x, dataset['movetext'])))
+else:
+    content = og_samples
+
+print('Data loaded')
+print('Training on ', len(content), 'characters. Good luck!')
 
 ## BUILD DATA SET ##
+# load data
+#with open('data.txt', 'r') as f:
+#    content = f.read()
+
 book = content
-characters = sorted(list(set(book)))
+if mask_all_data:
+    characters = sorted(list(set(book)))
+else:
+    characters = sorted(list(set('\n'.join(book))))
 vocab_size = len(characters)
 
 # convert 
-stoi = {ch: idx for idx, ch in enumerate(characters)}
-itos = {idx: ch for idx, ch in enumerate(characters)}
+class Tokenizer:
+    def __init__(self, vocab):
+        self.vocab = vocab
+        self.stoi = {ch: idx for idx, ch in enumerate(vocab)}
+        self.itos = {idx: ch for idx, ch in enumerate(vocab)}
+
+    def encode(self, s):
+        return [self.stoi[c] for c in s]
 
-encode = lambda s: [stoi[c] for c in s]
-decode = lambda i: ''.join([itos[x] for x in i])
+    def decode(self, i):
+        return ''.join([self.itos[x] for x in i])
 
+    @classmethod
+    def from_pretrained(cls, path):
+        with open(path, 'r') as f:
+            vocab = json.load(f)
+        return cls(vocab)
 
-data = torch.tensor(encode(book), dtype=torch.long)
-n = int(0.9 * len(data))
+    def save_pretrained(self, path):
+        with open(path, 'w') as f:
+            json.dump(self.vocab, f)
+
+
+tokenizer = Tokenizer(characters)
+encode = tokenizer.encode
+decode = tokenizer.decode
+
+if mask_all_data:
+    data = torch.tensor(encode(book), dtype=torch.long)
+else:
+    data = [torch.tensor(encode(s), dtype=torch.long) for s in book]
+max_len = max(len(x) for x in og_samples)
+context_size = min(context_size, max_len)
+
+
+n = int(0.8 * len(data))
 train_data = data[:n]
 val_data = data[n:]
 
 
+
+# Constants for piece movement validation
+PIECE_VALUES = {
+    'P': 1, 'N': 3, 'B': 3, 'R': 5, 'Q': 9, 'K': 0,  # White pieces
+    'p': 1, 'n': 3, 'b': 3, 'r': 5, 'q': 9, 'k': 0   # Black pieces
+}
+
+def initialize_board():
+    """Initializes the standard chessboard setup."""
+    return [
+        ['r', 'n', 'b', 'q', 'k', 'b', 'n', 'r'],  # 8th rank (Black)
+        ['p', 'p', 'p', 'p', 'p', 'p', 'p', 'p'],  # 7th rank (Black)
+        ['.', '.', '.', '.', '.', '.', '.', '.'],  # 6th rank
+        ['.', '.', '.', '.', '.', '.', '.', '.'],  # 5th rank
+        ['.', '.', '.', '.', '.', '.', '.', '.'],  # 4th rank
+        ['.', '.', '.', '.', '.', '.', '.', '.'],  # 3rd rank
+        ['P', 'P', 'P', 'P', 'P', 'P', 'P', 'P'],  # 2nd rank (White)
+        ['R', 'N', 'B', 'Q', 'K', 'B', 'N', 'R']   # 1st rank (White)
+    ]
+
+def get_piece(board, position):
+    """Returns the piece at a given board position (e.g., e4 -> 'P' or '.')."""
+    col = ord(position[0]) - ord('a')
+    row = 8 - int(position[1])
+    return board[row][col]
+
+def set_piece(board, position, piece):
+    """Sets a piece on the board at a given position."""
+    col = ord(position[0]) - ord('a')
+    row = 8 - int(position[1])
+    board[row][col] = piece
+
+def validate_pawn_move(board, start, end, is_white_turn):
+    """Validates pawn movement including capturing, advancing, and promotion."""
+    start_col, start_row = ord(start[0]) - ord('a'), 8 - int(start[1])
+    end_col, end_row = ord(end[0]) - ord('a'), 8 - int(end[1])
+    
+    pawn_direction = -1 if is_white_turn else 1  # White moves up, black moves down
+    
+    # Regular forward move
+    if start_col == end_col and board[end_row][end_col] == '.':
+        if start_row + pawn_direction == end_row:  # 1 square move
+            return True
+        if (is_white_turn and start_row == 6 or not is_white_turn and start_row == 1) and start_row + 2 * pawn_direction == end_row:
+            return True
+    
+    # Capture
+    if abs(start_col - end_col) == 1 and start_row + pawn_direction == end_row:
+        target_piece = board[end_row][end_col]
+        if (is_white_turn and target_piece.islower()) or (not is_white_turn and target_piece.isupper()):
+            return True
+    
+    return False
+
+def validate_knight_move(start, end):
+    """Validates knight movement (L-shape)."""
+    start_col, start_row = ord(start[0]) - ord('a'), 8 - int(start[1])
+    end_col, end_row = ord(end[0]) - ord('a'), 8 - int(end[1])
+    
+    col_diff = abs(start_col - end_col)
+    row_diff = abs(start_row - end_row)
+    
+    return (col_diff == 2 and row_diff == 1) or (col_diff == 1 and row_diff == 2)
+
+def validate_rook_move(board, start, end):
+    """Validates rook movement (straight lines along rank or file)."""
+    start_col, start_row = ord(start[0]) - ord('a'), 8 - int(start[1])
+    end_col, end_row = ord(end[0]) - ord('a'), 8 - int(end[1])
+    
+    if start_col != end_col and start_row != end_row:
+        return False  # Must be either same column or row
+    
+    # Check if path is clear
+    if start_col == end_col:
+        step = 1 if end_row > start_row else -1
+        for row in range(start_row + step, end_row, step):
+            if board[row][start_col] != '.':
+                return False
+    else:
+        step = 1 if end_col > start_col else -1
+        for col in range(start_col + step, end_col, step):
+            if board[start_row][col] != '.':
+                return False
+    
+    return True
+
+def validate_bishop_move(board, start, end):
+    """Validates bishop movement (diagonals)."""
+    start_col, start_row = ord(start[0]) - ord('a'), 8 - int(start[1])
+    end_col, end_row = ord(end[0]) - ord('a'), 8 - int(end[1])
+    
+    if abs(start_col - end_col) != abs(start_row - end_row):
+        return False  # Must move diagonally
+    
+    # Check if path is clear
+    col_step = 1 if end_col > start_col else -1
+    row_step = 1 if end_row > start_row else -1
+    col, row = start_col + col_step, start_row + row_step
+    while col != end_col and row != end_row:
+        if board[row][col] != '.':
+            return False
+        col += col_step
+        row += row_step
+    
+    return True
+
+def validate_move(board, move, is_white_turn):
+    """Validates a move based on the current board state."""
+    if move == "O-O" or move == "O-O-O":
+        return True  # Castling placeholder
+    
+    piece_type = 'P' if move[0].islower() else move[0]
+    start = move[-2:]  # Simplification; would need to parse actual source square
+    end = move[-2:]  # Actual end position is the destination
+    
+    if piece_type == 'P':
+        return validate_pawn_move(board, start, end, is_white_turn)
+    elif piece_type == 'N':
+        return validate_knight_move(start, end)
+    elif piece_type == 'R':
+        return validate_rook_move(board, start, end)
+    elif piece_type == 'B':
+        return validate_bishop_move(board, start, end)
+    
+    # Other pieces can be added similarly
+    return True  # Placeholder for other pieces
+
+def update_board(board, move, is_white_turn):
+    """Updates the board according to the move."""
+    start = move[-2:]
+    end = move[-2:]
+    piece = get_piece(board, start)
+    
+    # Move the piece
+    set_piece(board, end, piece)
+    set_piece(board, start, '.')
+    
+    return board  # Placeholder for now
+
+def validate_pgn(pgn_string):
+    """
+    Validates the PGN string format and chess move legality.
+    """
+    
+    move_pattern = r'([PNBRQK]?[a-h]?[1-8]?[x]?[a-h][1-8](=[QRNB])?|O-O(-O)?)[+#]?'  # Chess move
+    result_pattern = r'(1-0|0-1|1/2-1/2)'  # Game results
+    tag_pattern = r'\[([A-Za-z0-9_]+)\s+"([^"]+)"\]'  # PGN tags
+
+    pgn_lines = pgn_string.strip().splitlines()
+    
+    tags = [line for line in pgn_lines if line.startswith('[')]
+    for tag in tags:
+        if not re.match(tag_pattern, tag):
+            return False  # Invalid tag format
+    
+    moves_section = ' '.join([line for line in pgn_lines if not line.startswith('[')]).strip()
+    
+    if not re.search(result_pattern, moves_section):
+        return False  # No valid result found
+    
+    moves_section = re.sub(result_pattern, '', moves_section).strip()
+
+    board = initialize_board()
+    is_white_turn = True
+
+    move_tokens = re.split(r'\s|\d+\.', moves_section)
+    for token in move_tokens:
+        if token:
+            if not re.match(move_pattern, token):
+                return False  # Invalid move format
+            
+            if not validate_move(board, token, is_white_turn):
+                return False  # Invalid chess move
+            
+            board = update_board(board, token, is_white_turn)
+            is_white_turn = not is_white_turn
+    
+    return True
+
+# Test case
+pgn_string = """
+[Event "World Championship"]
+[Site "Moscow URS"]
+[Date "1985.11.09"]
+[Round "16"]
+[White "Kasparov, Garry"]
+[Black "Karpov, Anatoly"]
+[Result "1-0"]
+
+1. e4 e5 2. Nf3 Nc6 3. Bb5 a6 4. Ba4 Nf6 5. O-O Be7 6. Re1 b5 7. Bb3 d6
+8. c3 O-O 9. h3 Nb8 10. d4 Nbd7 11. c4 Bb7 12. Nbd2 c6 13. Bc2 Re8 14. b3 Bf8
+15. Bb2 Qc7 16. Rc1 Rad8 17. a3 Qb8 18. Bd3 g6 19. Qc2 Nh5 20. g3 Ng7 21. Qb1
+exd4 22. Nxd4 c5 23. N4f3 Ne6 24. Bf1 Ne5 25. Qa1 Nxf3+ 26. Nxf3 Qa8 27. b4
+Rc8 28. Bd3 Bh6 29. Rc2 Bc6 30. h4 f5 31. exf5 Bxf3 32. fxe6 Bh1 33. Bf1 Qf3
+34. Re2 Bg7 35. Kh2 Rc7 36. Bxg7 Rxg7 37. Qf6 bxc4 38. e7 Qxf6 39. exf6 1-0
+"""
+
+
+
+def get_batch_from_samples(split):
+    data = train_data if split == 'train' else val_data
+    sample_idx = torch.randint(len(data), (batch_size,))
+    inputs = []
+    outputs = []
+    space = encode(' ')[0]
+    for idx in sample_idx:
+        sample_size = len(data[idx])
+        start = torch.randint(max(sample_size - 2, sample_size - context_size), (1,))
+        end = start + context_size
+        i1 = data[idx][start:end].tolist()
+        i2 = [space] * (context_size - len(i1))
+        input_sample = torch.tensor(i1 + i2)
+        o1 = data[idx][start+1:end+1].tolist()
+        o2 = [space] * (context_size - len(o1))
+        output_sample = torch.tensor(o1 + o2)
+
+        inputs.append(input_sample)
+        outputs.append(output_sample)
+
+    x = torch.stack(inputs)
+    y = torch.stack(outputs)
+    return x.to(device), y.to(device)
+
+
 def get_batch(split):
     data = train_data if split == 'train' else val_data
     idx = torch.randint(len(data) - context_size, (batch_size,))
@@ -50,6 +332,9 @@ def get_batch(split):
     y = torch.stack([data[i+1:i+context_size+1] for i in idx])
     return x.to(device), y.to(device)
 
+if not mask_all_data:
+    get_batch = get_batch_from_samples
+
 ## END BUILD DATA SET ##
 ## MODEL DEFINITION ##
 
@@ -72,15 +357,58 @@ def estimate_loss():
         for k in range(eval_iters):
             X, Y = get_batch(split)
             logits, loss = model(X, Y)
+            """
+            input_string = X[0].tolist()
+            gen = model.generate(X[0].view(1, -1), max_new_tokens=5, context_size=context_size)
+            o = tokenizer.decode(gen[0].tolist())
+            try:
+                valid = int(not validate_pgn(o))
+            except Exception:
+                valid = 2
+            """
             losses[k] = loss.item()
         out[split] = losses.mean()
 
-    input_string = '1. e4 g6'
+    input_string = '1. e4 g6 2.'
     print_sample(torch.tensor(encode(input_string), dtype=torch.long, device=device).view((1, len(input_string))))
     model.train()
     return out
 
 
+class CosineAnnealingScheduler(_LRScheduler):
+    def __init__(self, optimizer, T_max, eta_min=0, last_epoch=-1):
+        """
+        Args:
+            optimizer (Optimizer): Wrapped optimizer.
+            T_max (int): Maximum number of iterations.
+            eta_min (float): Minimum learning rate. Default: 0.
+            last_epoch (int): The index of last epoch. Default: -1.
+        """
+        self.T_max = T_max
+        self.eta_min = eta_min
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if not self._get_lr_called_within_step:
+            warnings.warn("To get the last learning rate computed by the scheduler, "
+                          "please use `get_last_lr()`.", UserWarning)
+
+        if self.last_epoch == 0:
+            return [group['lr'] for group in self.optimizer.param_groups]
+        elif self._step_count == 1 and self.last_epoch > 0:
+            return [self.eta_min + (base_lr - self.eta_min) *
+                    (1 + math.cos((self.last_epoch) * math.pi / self.T_max)) / 2
+                    for base_lr in self.base_lrs]
+        elif (self.last_epoch - 1 - self.T_max) % (2 * self.T_max) == 0:
+            return [group['lr'] + (base_lr - self.eta_min) *
+                    (1 - math.cos(math.pi / self.T_max)) / 2
+                    for base_lr, group in
+                    zip(self.base_lrs, self.optimizer.param_groups)]
+        return [(1 + math.cos(math.pi * self.last_epoch / self.T_max)) /
+                (1 + math.cos(math.pi * (self.last_epoch - 1) / self.T_max)) *
+                (group['lr'] - self.eta_min) + self.eta_min
+                for group in self.optimizer.param_groups]
+
 if __name__ == "__main__":
     args = argparse.ArgumentParser()
     args.add_argument('--load', '-l', action='store_true', default=False, help='Load model state.')
@@ -91,28 +419,62 @@ if __name__ == "__main__":
     params = {'vocab_size': vocab_size, 'n_embed': n_embed, 'context_size': context_size, 'n_layer': n_layer, 'n_head': n_head, 'dropout': dropout}
     if args.load:
         m = DecoderTransformer(vocab_size, n_embed, context_size, n_layer, n_head, dropout)
-        m.load_state_dict(torch.load(f'./models/{base_name}' + ''.join(f'{key}={v}' for key, v in params.items())))
+        m.load_state_dict(torch.load(f'./models/{base_name}'))# + ''.join(f'{key}={v}' for key, v in params.items())))
     else:
         m = DecoderTransformer(vocab_size, n_embed, context_size, n_layer, n_head, dropout)
     model = m.to(device)
 
     if args.inference:
+        input_string = input('Enter a PGN string: ')
+        print_sample(torch.tensor(encode(input_string), dtype=torch.long, device=device).view((1, len(input_string))))
+        with open(f'./models/{base_name}_params.json', 'w') as f:
+            json.dump(params, f)
+
+        tokenizer.save_pretrained(f'./models/{base_name}_vocab.json')
         exit()
     ## END MODEL ##
     ## START TRAINING ##
+    wandb.init(project='chessPT')
+
+    wandb.watch(model)
     optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+    if use_scheduler:
+        scheduler = CosineAnnealingScheduler(optimizer, max_iters, eta_min=learning_rate//1e6)
 
-    for step in range(max_iters):
+    for step in tqdm(range(max_iters), total=max_iters, desc='Training'):
         if step % eval_interval == 0:
             losses = estimate_loss()
-            print(f'step {step:4d}: train loss {losses["train"]:.4f}, val loss: {losses["val"]:.4f}')
+            if use_scheduler:
+                print(f'step {step:4d}: train loss {losses["train"]:.4f}, val loss: {losses["val"]:.4f}, lr: {scheduler.get_last_lr()[0]}')
+            else:
+                print(f'step {step:4d}: train loss {losses["train"]:.4f}, val loss: {losses["val"]:.4f}')
+            wandb.log({'train_loss': losses['train'], 'val_loss': losses['val']})
 
         xb, yb = get_batch('train')
 
         logits, loss = model(xb, yb)
+        """
+
+        input_string = xb[0].tolist()
+        gen = model.generate(xb[0].view(1, -1), max_new_tokens=5, context_size=context_size)
+        out = tokenizer.decode(gen[0].tolist())
+        try:
+            valid = int(not validate_pgn(out))
+        except Exception:
+            valid = 2
+        loss += valid
+        """
+
+        if use_scheduler:
+            wandb.log({'running_train_loss': loss.item(), 'lr': scheduler.get_last_lr()[0]})
+        else:
+            wandb.log({'running_train_loss': loss.item()})
+
         optimizer.zero_grad(set_to_none=True)
         loss.backward()
         optimizer.step()
+        if use_scheduler:
+            scheduler.step()
 
     print()
     print('Loss:')
@@ -124,6 +486,8 @@ if __name__ == "__main__":
     ## END VALIDATION ##
 
     # save model weights
-    torch.save(model.state_dict(), f'./models/{base_name}' + ''.join([f'{key}={v}' for key, v in params.items()]))
+    torch.save(model.state_dict(), f'./models/{base_name}')
+    with open(f'./models/{base_name}_params.json', 'w') as f:
+        json.dump(params, f)
     with open('train.log', 'a') as f:
         f.write(f'{max_iters},{learning_rate}\n')