Silas Blanchard commited on
Commit
596285c
·
unverified ·
1 Parent(s): 475709b

Create model-runner.py (Our model does work on hugging face!)

Browse files
Files changed (1) hide show
  1. src/model-runner.py +117 -0
src/model-runner.py ADDED
@@ -0,0 +1,117 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import chess
2
+ import random
3
+
4
+ def cleanup_output(text, prompt):
5
+ section = text[len(prompt):len(prompt) + 6]
6
+ #print("Section: %s" % section)
7
+ valid_letters = ['A','a','B','b','C','c','D','d','E','e','F','f','G','g','H','h']
8
+ valid_pieces = ['p','P','k','K','q','Q','r','R','b','B', 'n', 'N']
9
+ valid_numbers = ['1','2','3','4','5','6','7','8']
10
+
11
+ #if there are any syntatically moves in this string for pieces
12
+ countr = 0
13
+ while countr < 4:
14
+ if(section[countr] in valid_pieces and section[countr + 1] in valid_letters and section[countr + 2] in valid_numbers):
15
+ #print(section[countr:countr+3])
16
+ return ' ' + section[countr:countr+3]
17
+ countr+=1
18
+
19
+ #same as moves but for pawns
20
+ countr = 0
21
+ while countr < 5:
22
+ if(section[countr] in valid_letters and section[countr+1] in valid_numbers):
23
+ #print(section[countr:countr+2])
24
+ return ' ' + section[countr:countr+2]
25
+ countr+=1
26
+
27
+ return ' e8'
28
+
29
+ class AInstance:
30
+ def __init__(self, type, generator):
31
+ self.type = type
32
+ self.game_end = False
33
+ self.generator = generator
34
+
35
+ #All this does it take the gamestate and add the ai-generated result to it
36
+ def move(self, game_state):
37
+ if(type == "BlueSunflower/gpt2-medium-chess"):
38
+ prompt = "1-0 2700 1350 " + game_state
39
+ else:
40
+ prompt = game_state
41
+ countr = 0
42
+ while True:
43
+ generated_text = self.generator(prompt, max_length=len(prompt) + 10, num_return_sequences=1)[0]['generated_text']
44
+ selected_move = cleanup_output(generated_text, prompt)
45
+
46
+ #if this move is valid then return it
47
+ proposed_board = game_state + selected_move
48
+ if(verify_move(proposed_board)):
49
+ return proposed_board
50
+ countr+=1
51
+ #goes fifty times until the AInstance object flags itself as "ended" (fundamentally unable to make a valid move)
52
+ if(countr > 50):
53
+ self.game_end = True
54
+ break
55
+
56
+ def check_if_end(self):
57
+ return self.game_end
58
+
59
+ def verify_move(string):
60
+ board = chess.Board()
61
+ print("Board: %s" % string)
62
+ for move in string.split():
63
+ #if this move makes no sense it will return false and the game will try again to generate a good move
64
+ try:
65
+ board.push_san(move)
66
+ except:
67
+ return False
68
+ if(board.is_valid):
69
+ return True
70
+ return False
71
+
72
+ def check_mate(string):
73
+ #simulates mate idk
74
+ if(random.randrange(0,100) == 4):
75
+ print("H")
76
+ return True
77
+ return False
78
+
79
+ def print_game(string):
80
+ print("Some kind of visualization for the chess board based on this string: %s" % string)
81
+
82
+ def make_move(instance, game_state):
83
+ print("%s's move" % instance.type)
84
+ return_state = game_state
85
+ return_state = instance.move(game_state)
86
+ game_ongoing = True
87
+ if(instance.check_if_end()):
88
+ print("This player claims countr > 50: %s" % instance.type)
89
+ game_ongoing = False
90
+ if(check_mate(return_state)):
91
+ print("This player claims mates: %s" % instance.type)
92
+ game_ongoing = False
93
+ return(return_state, game_ongoing)
94
+
95
+
96
+ def main():
97
+ if(random.randrange(0,1)):
98
+ white = AInstance("gpt2", generator)
99
+ black = AInstance("gpt2-medium-chess", generator2)
100
+ print("Gpt2 is White and Gpt2 Optimized is Black")
101
+ else:
102
+ white = AInstance("gpt2-medium-chess", generator2)
103
+ black = AInstance("gpt2", generator)
104
+ print("Gpt2 is Black and Gpt2 Optimized is White")
105
+
106
+ game_state = "e4 e5"
107
+ game_ongoing = True
108
+ while game_ongoing:
109
+ game_state, game_ongoing = make_move(white, game_state)
110
+ if not game_ongoing:
111
+ print_game(game_state)
112
+ break
113
+ game_state, game_ongoing = make_move(black, game_state)
114
+ if not game_ongoing:
115
+ print_game(game_state)
116
+ break
117
+ main()