File size: 3,239 Bytes
1408cee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | import berserk
import torch
from torch import nn
from torch import optim
import chess
import os
import pandas as pd
device = torch.device("cpu")
def board_to_tensor(board):
piece_encoding = {
'P': 1, 'N': 2, 'B': 3, 'R': 4, 'Q': 5, 'K': 6,
'p': 7, 'n': 8, 'b': 9, 'r': 10, 'q': 11, 'k': 12
}
tensor = torch.zeros(64, dtype=torch.long)
for square in chess.SQUARES:
piece = board.piece_at(square)
if piece:
tensor[square] = piece_encoding[piece.symbol()]
else:
tensor[square] = 0
return tensor.unsqueeze(0)
class BOT(nn.Module):
def __init__(self):
super().__init__()
self.embedding = nn.Embedding(13,64)
self.attention = nn.MultiheadAttention(embed_dim=64,num_heads=16)
self.neurons = nn.Sequential(
nn.Linear(4096,128),
nn.ReLU(),
nn.Linear(128,128),
nn.ReLU(),
nn.Linear(128,128),
nn.ReLU(),
nn.Linear(128,64),
nn.ReLU(),
nn.Linear(64,1)
)
def forward(self,x):
x = self.embedding(x)
x = x.permute(1, 0, 2)
attn_output, _ = self.attention(x, x, x)
x = attn_output.permute(1, 0, 2).contiguous()
x = x.view(x.size(0), -1)
x = self.neurons(x)
return x
model = BOT().to(device)
model = torch.compile(model,mode="max-autotune",dynamic=False)
if os.path.exists("booty.pth"):
file = torch.load("booty.pth",map_location=device,weights_only=True)
model.load_state_dict(file)
model.train()
optimizer = optim.Adam(model.parameters(),lr=1e-4)
criterion = nn.MSELoss()
num_epochs = 1
df = pd.read_csv("lichess_db_puzzle.csv",nrows=1000)
df = df.sort_values(by="Rating", ascending=True)
t1 = torch.tensor([10.0], dtype=torch.float32, device=device)
t2 = torch.tensor([-10.0], dtype=torch.float32, device=device)
t3 = torch.tensor([0.0], dtype=torch.float32, device=device)
for i in range(num_epochs):
total_loss = 0.0
for puzzle in df.iloc:
board = chess.Board(puzzle["FEN"])
print(f"Rating: {puzzle['Rating']} elo.")
n = 0
for move in puzzle["Moves"].split():
if n % 2 == 0:
for movey in list(board.legal_moves):
if str(movey.uci()) == move:
b = True
else:
b = False
board.push(movey)
tensor = board_to_tensor(board).to(device)
evaling = model(tensor)
board.pop()
optimizer.zero_grad()
if b and board.turn == chess.WHITE:
loss = criterion(t1, evaling)
elif b and board.turn == chess.BLACK:
loss = criterion(t2, evaling)
else:
loss = criterion(t3, evaling)
total_loss += loss.item()
loss.backward()
optimizer.step()
n += 1
board.push_uci(move)
print(f"Epoch [{i+1}/{num_epochs}], Loss: {total_loss}.")
torch.save(model.state_dict(),"booty.pth")
|