initial commit

ai.py

@@ -0,0 +1,39 @@
+import numpy as np
+import torch
+from transformers import GPT2LMHeadModel
+
+
+def load_model(model_name: str = "snoop2head/Gomoku-GPT2") -> GPT2LMHeadModel:
+    gpt2 = GPT2LMHeadModel.from_pretrained(model_name)
+    return gpt2
+
+
+BOS_TOKEN_ID = 401
+PAD_TOKEN_ID = 402
+EOS_TOKEN_ID = 403
+
+
+def generate_gpt2(model: GPT2LMHeadModel, input_ids: torch.LongTensor) -> list:
+    """
+    input_ids: [batch_size, seq_len] torch.LongTensor
+    output_ids: [seq_len] list
+    """
+    output_ids = model.generate(
+        input_ids,
+        max_length=128,
+        num_beams=5,
+        temperature=0.7,
+        pad_token_id=PAD_TOKEN_ID,
+        eos_token_id=EOS_TOKEN_ID,
+    )
+    return output_ids.squeeze().tolist()
+
+
+def change_to_1d_coordinate(board: np.ndarray, x: int, y: int) -> int:
+    """change 2d coordinate to 1d coordinate"""
+    return x * board.shape[1] + y
+
+
+def change_to_2d_coordinate(board: np.ndarray, coordinate: int) -> tuple:
+    """change 1d coordinate to 2d coordinate"""
+    return (coordinate // board.shape[1], coordinate % board.shape[1])

app.py

@@ -0,0 +1,399 @@
+"""
+- This is a simple gomoku game built with Streamlit by TeddyHuang-00 (huang_nan_2019@pku.edu.cn).
+- For Gomoku-GPT2, please refer to Young-Jin Ahn (young_ahn@yonsei.ac.kr).
+
+Shared under MIT license
+"""
+
+import time
+from copy import deepcopy
+from uuid import uuid4
+
+import torch
+import numpy as np
+import streamlit as st
+from scipy.signal import convolve
+from streamlit import session_state
+from streamlit_server_state import server_state, server_state_lock
+
+from ai import (
+    BOS_TOKEN_ID,
+    generate_gpt2,
+    load_model,
+)
+
+
+# Utils
+class Room:
+    def __init__(self, room_id) -> None:
+        self.ROOM_ID = room_id
+        self.BOARD = np.zeros(shape=(20, 20), dtype=int)
+        self.PLAYER = _BLACK
+        self.TURN = self.PLAYER
+        self.HISTORY = (0, 0)
+        self.WINNER = _BLANK
+        self.TIME = time.time()
+        self.COORDINATE_1D = [BOS_TOKEN_ID]
+
+
+gpt2 = load_model()
+
+
+_BLANK = 0
+_BLACK = 1
+_WHITE = -1
+_PLAYER_SYMBOL = {
+    _WHITE: "⚪",
+    _BLANK: "➕",
+    _BLACK: "⚫",
+}
+_PLAYER_COLOR = {
+    _WHITE: "Gomoku-GPT",
+    _BLANK: "Blank",
+    _BLACK: "YOU HUMAN",
+}
+_HORIZONTAL = np.array(
+    [
+        [0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0],
+        [1, 1, 1, 1, 1],
+        [0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0],
+    ]
+)
+_VERTICAL = np.array(
+    [
+        [0, 0, 1, 0, 0],
+        [0, 0, 1, 0, 0],
+        [0, 0, 1, 0, 0],
+        [0, 0, 1, 0, 0],
+        [0, 0, 1, 0, 0],
+    ]
+)
+_DIAGONAL_UP_LEFT = np.array(
+    [
+        [1, 0, 0, 0, 0],
+        [0, 1, 0, 0, 0],
+        [0, 0, 1, 0, 0],
+        [0, 0, 0, 1, 0],
+        [0, 0, 0, 0, 1],
+    ]
+)
+_DIAGONAL_UP_RIGHT = np.array(
+    [
+        [0, 0, 0, 0, 1],
+        [0, 0, 0, 1, 0],
+        [0, 0, 1, 0, 0],
+        [0, 1, 0, 0, 0],
+        [1, 0, 0, 0, 0],
+    ]
+)
+
+_ROOM_COLOR = {
+    True: _BLACK,
+    False: _WHITE,
+}
+
+# Initialize the game
+if "ROOM" not in session_state:
+    session_state.ROOM = Room("local")
+if "OWNER" not in session_state:
+    session_state.OWNER = False
+
+# Check server health
+if "ROOMS" not in server_state:
+    with server_state_lock["ROOMS"]:
+        server_state.ROOMS = {}
+
+# # Layout
+# Main
+TITLE = st.empty()
+ROUND_INFO = st.empty()
+BOARD_PLATE = [
+    [cell.empty() for cell in st.columns([1 for _ in range(20)])] for _ in range(20)
+]
+WAIT_FOR_OPPONENT = st.empty()
+
+# Sidebar
+SCORE_TAG = st.sidebar.empty()
+SCORE_PLATE = st.sidebar.columns(2)
+PLAY_MODE_INFO = st.sidebar.container()
+MULTIPLAYER_TAG = st.sidebar.empty()
+with st.sidebar.container():
+    ANOTHER_ROUND = st.empty()
+    RESTART = st.empty()
+    EXIT = st.empty()
+GAME_INFO = st.sidebar.container()
+
+
+# Draw the board
+def gomoku():
+    """
+    Draw the board.
+
+    Handle the main logic.
+    """
+
+    # Restart the game
+    def restart() -> None:
+        """
+        Restart the game.
+        """
+        session_state.ROOM = Room(session_state.ROOM.ROOM_ID)
+
+    # Continue new round
+    def another_round() -> None:
+        """
+        Continue new round.
+        """
+        session_state.ROOM = deepcopy(session_state.ROOM)
+        session_state.ROOM.BOARD = np.zeros(shape=(20, 20), dtype=int)
+        session_state.ROOM.PLAYER = -session_state.ROOM.PLAYER
+        session_state.ROOM.TURN = session_state.ROOM.PLAYER
+        session_state.ROOM.WINNER = _BLANK
+        session_state.ROOM.COORDINATE_1D = [BOS_TOKEN_ID]
+
+    # Room status sync
+    def sync_room() -> bool:
+        room_id = session_state.ROOM.ROOM_ID
+        if room_id not in server_state.ROOMS.keys():
+            session_state.ROOM = Room("local")
+            return False
+        elif server_state.ROOMS[room_id].TIME == session_state.ROOM.TIME:
+            return False
+        elif server_state.ROOMS[room_id].TIME < session_state.ROOM.TIME:
+            # Only acquire the lock when writing to the server state
+            with server_state_lock["ROOMS"]:
+                server_rooms = server_state.ROOMS
+                server_rooms[room_id] = session_state.ROOM
+                server_state.ROOMS = server_rooms
+            return True
+        else:
+            session_state.ROOM = server_state.ROOMS[room_id]
+            return True
+
+    # Check if winner emerge from move
+    def check_win() -> int:
+        """
+        Use convolution to check if any player wins.
+        """
+        vertical = convolve(
+            session_state.ROOM.BOARD,
+            _VERTICAL,
+            mode="same",
+        )
+        horizontal = convolve(
+            session_state.ROOM.BOARD,
+            _HORIZONTAL,
+            mode="same",
+        )
+        diagonal_up_left = convolve(
+            session_state.ROOM.BOARD,
+            _DIAGONAL_UP_LEFT,
+            mode="same",
+        )
+        diagonal_up_right = convolve(
+            session_state.ROOM.BOARD,
+            _DIAGONAL_UP_RIGHT,
+            mode="same",
+        )
+        if (
+            np.max(
+                [
+                    np.max(vertical),
+                    np.max(horizontal),
+                    np.max(diagonal_up_left),
+                    np.max(diagonal_up_right),
+                ]
+            )
+            == 5 * _BLACK
+        ):
+            winner = _BLACK
+        elif (
+            np.min(
+                [
+                    np.min(vertical),
+                    np.min(horizontal),
+                    np.min(diagonal_up_left),
+                    np.min(diagonal_up_right),
+                ]
+            )
+            == 5 * _WHITE
+        ):
+            winner = _WHITE
+        else:
+            winner = _BLANK
+        return winner
+
+    # Triggers the board response on click
+    def handle_click(x, y):
+        """
+        Controls whether to pass on / continue current board / may start new round
+        """
+        if session_state.ROOM.BOARD[x][y] != _BLANK:
+            pass
+        elif (
+            session_state.ROOM.ROOM_ID in server_state.ROOMS.keys()
+            and _ROOM_COLOR[session_state.OWNER]
+            != server_state.ROOMS[session_state.ROOM.ROOM_ID].TURN
+        ):
+            sync_room()
+
+        # normal play situation
+        elif session_state.ROOM.WINNER == _BLANK:
+            session_state.ROOM = deepcopy(session_state.ROOM)
+
+            session_state.ROOM.BOARD[x][y] = session_state.ROOM.TURN
+            session_state.ROOM.COORDINATE_1D.append(x * 20 + y)
+
+            session_state.ROOM.TURN = -session_state.ROOM.TURN
+            session_state.ROOM.WINNER = check_win()
+            session_state.ROOM.HISTORY = (
+                session_state.ROOM.HISTORY[0]
+                + int(session_state.ROOM.WINNER == _WHITE),
+                session_state.ROOM.HISTORY[1]
+                + int(session_state.ROOM.WINNER == _BLACK),
+            )
+            session_state.ROOM.TIME = time.time()
+
+    # Draw board
+    def draw_board(response: bool):
+        """construct each buttons for all cells of the board"""
+
+        if response and session_state.ROOM.TURN == 1:  # human turn
+            print("Your turn")
+            # construction of clickable buttons
+            for i, row in enumerate(session_state.ROOM.BOARD):
+                for j, cell in enumerate(row):
+                    BOARD_PLATE[i][j].button(
+                        _PLAYER_SYMBOL[cell],
+                        key=f"{i}:{j}",
+                        on_click=handle_click,
+                        args=(i, j),
+                    )
+
+        elif response and session_state.ROOM.TURN == -1:  # AI turn
+            print("AI's turn")
+            gpt_predictions = generate_gpt2(
+                gpt2,
+                torch.tensor(session_state.ROOM.COORDINATE_1D).unsqueeze(0),
+            )
+            print(gpt_predictions)
+            gpt_response = gpt_predictions[len(session_state.ROOM.COORDINATE_1D)]
+            gpt_i, gpt_j = gpt_response // 20, gpt_response % 20
+            print(gpt_i, gpt_j)
+            session_state.ROOM.BOARD[gpt_i][gpt_j] = session_state.ROOM.TURN
+            session_state.ROOM.COORDINATE_1D.append(gpt_i * 20 + gpt_j)
+
+            # construction of clickable buttons
+            for i, row in enumerate(session_state.ROOM.BOARD):
+                for j, cell in enumerate(row):
+                    if (
+                        i * 20 + j
+                        in gpt_predictions[: len(session_state.ROOM.COORDINATE_1D)]
+                    ):
+                        # disable click for GPT choices
+                        BOARD_PLATE[i][j].button(
+                            _PLAYER_SYMBOL[cell],
+                            key=f"{i}:{j}",
+                            on_click=False,
+                            args=(i, j),
+                        )
+                    else:
+                        # enable click for other cells available for human choices
+                        BOARD_PLATE[i][j].button(
+                            _PLAYER_SYMBOL[cell],
+                            key=f"{i}:{j}",
+                            on_click=handle_click,
+                            args=(i, j),
+                        )
+
+            # change turn
+            session_state.ROOM.TURN = -session_state.ROOM.TURN
+            session_state.ROOM.WINNER = check_win()
+            session_state.ROOM.HISTORY = (
+                session_state.ROOM.HISTORY[0]
+                + int(session_state.ROOM.WINNER == _WHITE),
+                session_state.ROOM.HISTORY[1]
+                + int(session_state.ROOM.WINNER == _BLACK),
+            )
+            session_state.ROOM.TIME = time.time()
+
+        if not response or session_state.ROOM.WINNER != _BLANK:
+            print("Game over")
+            for i, row in enumerate(session_state.ROOM.BOARD):
+                for j, cell in enumerate(row):
+                    BOARD_PLATE[i][j].write(
+                        _PLAYER_SYMBOL[cell],
+                        key=f"{i}:{j}",
+                    )
+
+    # Game process control
+    def game_control():
+        if session_state.ROOM.WINNER != _BLANK:
+            draw_board(False)
+        else:
+            draw_board(True)
+        if session_state.ROOM.WINNER != _BLANK or 0 not in session_state.ROOM.BOARD:
+            ANOTHER_ROUND.button(
+                "Play Next round!",
+                on_click=another_round,
+                help="Clear board and swap first player",
+            )
+        if session_state.ROOM.ROOM_ID == "local" or session_state.OWNER:
+            RESTART.button(
+                "Reset",
+                on_click=restart,
+                help="Clear the board as well as the scores",
+            )
+
+    # Infos
+    def draw_info() -> None:
+        # Text information
+        TITLE.subheader("**🤖 Do you wanna have a bad time?**")
+        PLAY_MODE_INFO.write("---\n\n**You are Black, AI is White.**")
+        GAME_INFO.markdown(
+            """
+            ---
+
+            ## Freestyle Gomoku game.
+
+
+            <a href="https://en.wikipedia.org/wiki/Gomoku#Freestyle_Gomoku" style="color:#FFFFFF">Freestyle Gomoku</a>
+
+            - no restrictions
+            - no regrets
+            - swap players after one round is over
+
+            ##### Design by <a href="https://github.com/TeddyHuang-00" style="color:#FFFFFF">TeddyHuang-00</a> • <a href="https://github.com/TeddyHuang-00/streamlit-gomoku" style="color:#FFFFFF">Github repo</a>
+            ##### Gomoku-GPT by <a href="https://github.com/snoop2head" style="color:#FFFFFF">snoop2head</a> • <a href="https://github.com/snoop2head/" style="color:#FFFFFF">Github repo</a>
+
+            """,
+            unsafe_allow_html=True,
+        )
+        # History scores
+        SCORE_TAG.subheader("Scores")
+        SCORE_PLATE[0].metric("Gomoku-GPT", session_state.ROOM.HISTORY[0])
+        SCORE_PLATE[1].metric("Black", session_state.ROOM.HISTORY[1])
+
+        # Additional information
+        if session_state.ROOM.WINNER != _BLANK:
+            st.balloons()
+            ROUND_INFO.write(
+                f"#### **{_PLAYER_COLOR[session_state.ROOM.WINNER]} wins!**\n**Click buttons on the left for more plays.**"
+            )
+
+        elif 0 not in session_state.ROOM.BOARD:
+            ROUND_INFO.write("#### **Tie**")
+        else:
+            ROUND_INFO.write(
+                f"#### **{_PLAYER_SYMBOL[session_state.ROOM.TURN]} {_PLAYER_COLOR[session_state.ROOM.TURN]}'s turn...**"
+            )
+
+    # The main game loop
+    game_control()
+    draw_info()
+
+
+if __name__ == "__main__":
+    gomoku()

requirements.txt

@@ -0,0 +1,6 @@
+numpy==1.23.5
+scipy==1.10.1
+torch==2.0.0
+streamlit==0.89.0
+streamlit_server_state==0.6.1
+transformers==4.27.4