Views
No views yet
git clone https://github.com/waterhorse1/ChessGPT1import sys
2sys.path.append('./chessclip/src')
3import torch
4import io
5import chess.pgn
6import numpy as np
7from chess_ai.feature_converter import get_lc0_input_planes_tf
8from chess_ai.datasets.tfds.pgn_base import generate_examples_from_game_no_comment
9
10from open_clip.factory import get_tokenizer, load_checkpoint
11
12# init
13model_name = 'chessclip-quickgelu'
14model = open_clip.create_model(model_name, pretrained='openai')
15tokenizer = get_tokenizer(model_name)
16
17# load model
18load_checkpoint(model, './ChessCLIP/epoch_latest.pt')
19
20# check parameters
21model.eval()
22context_length = model.text.context_length
23vocab_size = model.text.vocab_size
24
25print("Model parameters:", f"{np.sum([int(np.prod(p.shape)) for p in model.parameters()]):,}")
26print("Context length:", context_length)
27print("Vocab size:", vocab_size)
28
29# generate board/action embedding based on pgn string
30def generate_representation_for_final(pgn):
31 game = chess.pgn.read_game(io.StringIO(pgn))
32 data = list(generate_examples_from_game_no_comment(game))[-1]
33 for key in data.keys():
34 data[key] = np.array(data[key])
35 board = get_lc0_input_planes_tf(data).numpy()
36 action = data['probs']
37 return board, action
38
39# Prepare input
40prompt = "Black plays Sicilian Defense"
41pgn_str = '1. e4 c5'
42board, action = generate_representation_for_final(pgn_str)
43text_tokens = tokenizer([prompt])
44
45image_input = torch.from_numpy(np.stack([board], axis=0))
46action_input = torch.from_numpy(np.stack([action], axis=0))
47
48# infer
49with torch.no_grad():
50 image_features = model.encode_image((image_input, action_input)).float()
51 text_features = model.encode_text(text_tokens).float()
52image_features /= image_features.norm(dim=-1, keepdim=True) # n * dim
53text_features /= text_features.norm(dim=-1, keepdim=True)# m * dim
54similarity = text_features.cpu().numpy() @ image_features.cpu().numpy().T # m * n
55print(similarity)