Update app.py

app.py

@@ -3,21 +3,21 @@ import random
 import streamlit as st
 from transformers import pipeline
 
-def cleanup_output(text, prompt, extra_len):
-        section = text[len(prompt):]
-        st.write("Proposed Move: " + section)
+def cleanup_output(text, prompt):
+        section = text[len(prompt):len(prompt) + 7]
+        print("Proposed Move: %s" % section)
         valid_letters = ['A','a','B','b','C','c','D','d','E','e','F','f','G','g','H','h']
         valid_pieces = ['p','P','k','K','q','Q','r','R','b','B', 'n', 'N']
         valid_numbers = ['1','2','3','4','5','6','7','8']
 
         #if there are any syntatically moves in this string for pieces
         countr = 0
-        while countr < len(section) - 3:
+        while countr < 4:
           if(section[countr] in valid_pieces and section[countr + 1] in valid_letters and section[countr + 2] in valid_numbers):
             #print(section[countr:countr+3])
             return ' ' + section[countr:countr+3]
           countr+=1
-        
+
         #variant for capturing!
         countr = 0
         while countr < len(section) - 4:
@@ -28,12 +28,12 @@ def cleanup_output(text, prompt, extra_len):
 
         #same as moves but for pawns
         countr = 0
-        while countr < len(section) - 2:
+        while countr < 5:
           if(section[countr] in valid_letters and section[countr+1] in valid_numbers):
             #print(section[countr:countr+2])
             return ' ' + section[countr:countr+2]
           countr+=1
-
+        
         #variant for capturing!
         countr = 0
         while countr < len(section) -4:
@@ -41,7 +41,7 @@ def cleanup_output(text, prompt, extra_len):
             #print(section[countr:countr+2])
             return ' ' + section[countr:countr+4]
           countr+=1
-        
+
         return ' e8'
 
 class AInstance:
@@ -52,17 +52,15 @@ class AInstance:
 
     #All this does it take the gamestate and add the ai-generated result to it
     def move(self, game_state):
-        if(self.type == "gpt2-medium-chess"):
+        if(type == "BlueSunflower/gpt2-medium-chess"):
             prompt = "1-0 2700 1350 " + game_state
-            extra_len = 7
         else:
             prompt = game_state
-            extra_len = 5
         countr = 0
         while True:
-            generated_text = self.generator(prompt, max_length=len(prompt) + extra_len, num_return_sequences=1)[0]['generated_text']
-            selected_move = cleanup_output(generated_text, prompt, extra_len)
-            
+            generated_text = self.generator(prompt, max_length=len(prompt) + 10, num_return_sequences=1)[0]['generated_text']
+            selected_move = cleanup_output(generated_text, prompt)
+
             #if this move is valid then return it
             proposed_board = game_state + selected_move
             if(verify_move(proposed_board)):
@@ -78,53 +76,52 @@ class AInstance:
 
 def verify_move(string):
     board = chess.Board()
-    st.write("Board: " + string + "\n")
+    print("Board: %s" % string)
     for move in string.split():
       #if this move makes no sense it will return false and the game will try again to generate a good move
       try:
         board.push_san(move)
       except:
+        print("Invalid Move\n")
         return False
     if(board.is_valid):
+      print("Valid Move\n")
       return True
     return False
 
 def check_mate(string):
   #simulates mate idk
     if(random.randrange(0,100) == 4):
-        st.write("H")
+        print("H")
         return True
     return False
 
 def print_game(string):
-    st.write("Some kind of visualization for the chess board based on this string: " + string)
+    print("Some kind of visualization for the chess board based on this string: %s" % string)
 
 def make_move(instance, game_state):
-    st.write(instance.type + "'s move")
+    print("\n%s's move" % instance.type)
     return_state = game_state
     return_state = instance.move(game_state)
     game_ongoing = True
     if(instance.check_if_end()):
-        st.write("This player claims they can't make a valid move after 50 tries: " + instance.type)
+        print("This player claims countr > 50: %s" % instance.type)
         game_ongoing = False
     if(check_mate(return_state)):
-        st.write("This player claims mates: " +  instance.type)
+        print("This player claims mates: %s" % instance.type)
         game_ongoing = False
     return(return_state, game_ongoing)
 
 
 def main():
-    generator2 = pipeline('text-generation', model='BlueSunflower/gpt2-medium-chess')
-    generator = pipeline('text-generation', model='gpt2')
-    
-    if(random.randint(0,1) == 1):
+    if(random.randrange(0,1)):
         white = AInstance("gpt2", generator)
         black = AInstance("gpt2-medium-chess", generator2)
-        st.write("Gpt2 is White and Gpt2 Optimized is Black\n")
+        print("Gpt2 is White and Gpt2 Optimized is Black")
     else:
         white = AInstance("gpt2-medium-chess", generator2)
         black = AInstance("gpt2", generator)
-        st.write("Gpt2 is Black and Gpt2 Optimized is White\n")
+        print("Gpt2 is Black and Gpt2 Optimized is White")
 
     game_state = "e4 e5"
     game_ongoing = True
@@ -137,4 +134,4 @@ def main():
         if not game_ongoing:
             print_game(game_state)
             break
-main()
+main()