Update chess_board.py

chess_board.py

@@ -1,9 +1,9 @@
-# import os
-# os.environ["KERAS_BACKEND"] = "torch"  # "jax", "torch" or "tensorflow"
+import os
+os.environ["KERAS_BACKEND"] = "torch"  # "jax", "torch" or "tensorflow"
 
-# import keras_nlp
-# import keras
-# import torch
+import keras_nlp
+import keras
+import torch
 
 import chess
 import chess.svg
@@ -17,10 +17,10 @@ class Game:
         self.counter = 0
         self.arrow= None
         
-        # self.model_id = 'kaggle://valentinbaltazar/gemma-chess/keras/gemma_2b_en_chess'
-        # self.sampler = keras_nlp.samplers.TopKSampler(k=50, temperature=0.7)
-        # self.model = keras_nlp.models.GemmaCausalLM.from_preset(self.model_id)
-        # self.compile_model()
+        self.model_id = 'kaggle://valentinbaltazar/gemma-chess/keras/gemma_2b_en_chess'
+        self.sampler = keras_nlp.samplers.TopKSampler(k=50, temperature=0.7)
+        self.model = keras_nlp.models.GemmaCausalLM.from_preset(self.model_id)
+        self.compile_model()
 
     def compile_model(self):
         self.model.compile(sampler=self.sampler)
@@ -37,9 +37,9 @@ class Game:
                 instruction=f"Predict the next chess move in the sequence {str(self.sequence)}",
                 response="",)
 
-            # output = self.model.generate(prompt, max_length=256)
+            output = self.model.generate(prompt, max_length=256)
         
-            # gemma_move = output.split(' ')[-1].strip("'")
+            gemma_move = output.split(' ')[-1].strip("'")
 
         if self.make_move(gemma_move):
             print(f'Gemma plays {self.sequence[-1]}! (Current Sequence: {self.sequence} {len(self.sequence)})')
@@ -54,8 +54,7 @@ class Game:
             return None
 
     def gemma_moves(self):
-        # print(f"Gemma is thinking...(Current Sequence: {self.sequence} {len(self.sequence)})")
-        # time.sleep(3)
+        """Calls Gemma to make a move, either self generated or from opening sequence"""
         if self.opening_moves and len(self.sequence)<len(self.opening_moves):
             return self.call_gemma(self.opening_moves[len(self.sequence)])
         else:
@@ -64,23 +63,20 @@ class Game:
     def player_moves(self, move):
         return self.make_move(move)
 
-    # Function to display the board
     def display_board(self):
-        # clear_output(wait=True)
-        # display(SVG(chess.svg.board(board=self.board)))
+        """Return SVG image of board state"""
         if self.arrow:
             board_svg = chess.svg.board(board=self.board, arrows=[self.arrow])
         else:
             board_svg = chess.svg.board(board=self.board)
-        # return svg2png(bytestring=board_svg)
         return board_svg
 
-    # Function to make a move
+   
     def make_move(self, move):
+        """Checks to see if move is valid, if so pushes move to board state"""
         try:
             update = self.board.parse_san(move)
             self.board.push(update)
-            # self.display_board()
             self.sequence.append(move)
             self.arrow = chess.svg.Arrow(update.from_square, update.to_square, color="#0000cccc")
             return True
@@ -93,10 +89,10 @@ class Game:
         self.sequence = []
         self.counter = 0
         self.arrow = None
-        # self.board.reset
         return self.display_board()
     
     def generate_moves(self, move):
+        """Generator function for one full turn of chess moves"""
         valid_move = self.player_moves(move)
         if valid_move:
             yield self.display_board(), f"You played: {move}"