-
Notifications
You must be signed in to change notification settings - Fork 8
/
pgn_to_records.py
executable file
·127 lines (96 loc) · 3 KB
/
pgn_to_records.py
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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
#! /usr/bin/env python2.7
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from six.moves import xrange
import os
import sys
import chess
import chess.pgn
import numpy as np
import tensorflow as tf
import adapter
import output_fn
import self_play
import util
from config import config
def run_game (game):
# Initialize memory
actions = []
policies = []
indices = []
# Run through game
node = game
board = chess.Board()
while not node.is_end():
next_node = node.variation(0)
move = next_node.move
# Get action taken and action list
action = adapter.move_to_label_flat(move)
legal_actions = map(adapter.move_to_label_flat, board.legal_moves)
# Create one-hot probability vector
index = legal_actions.index(action)
probs = util.one_hot(index, len(legal_actions))
assert board.is_legal(move)
# TODO: Look at the validity of this in case of underpromotion
board.push(move)
node = next_node
# Update memory
actions.append(action)
policies.append(probs)
indices.append(legal_actions)
# Get game winner
winner, outcome = {
'1/2-1/2' : (chess.WHITE, 0.0),
'1-0' : (chess.WHITE, 1.0),
'0-1' : (chess.BLACK, 1.0)
}.get(game.headers['Result'], None)
return actions, policies, indices, outcome, winner
def run_pgn (pgn_file, n_games, data_dir):
games = 0
while n_games == 0 or games < n_games:
game = chess.pgn.read_game(pgn_file)
game.headers['Counter'] = games
name = '{White}-{Black}-{ECO}-{Date}-{Counter}'.format (
**game.headers
).replace(' ', '_')
# Loop exit condition
if game is None:
break
# Run through game generating labels
actions, policies, indices, outcome, winner = run_game(game)
# Save labels to disk
self_play.write_records(data_dir, name, actions, policies, indices, outcome, winner)
#
games += 1
def main (FLAGS, _):
# Parse pgn into registry of nodes
with open(FLAGS.pgn_file, 'r') as in_file:
run_pgn(in_file, FLAGS.n_games, FLAGS.data_path)
return 0
if __name__ == '__main__':
import argparse
import sys
parser = argparse.ArgumentParser()
parser.add_argument (
'--pgn_file', type=str, metavar='pgn',
required=True,
help='PGN database to parse.'
)
parser.add_argument (
'--data_path', type=str, metavar='dir',
required=True,
help='Where to store output data.'
)
parser.add_argument (
'--n_games', type=int, metavar='n',
default=0,
help='Number of games to parse.'
)
parser.add_argument (
'--temperature', type=float, metavar='T',
default=1.0,
help='Temperature to use in softcount.'
)
FLAGS, unknown = parser.parse_known_args()
exit(main(FLAGS, [sys.argv[0]] + unknown))