WordleAI_Test.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "WordleAI Test.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "authorship_tag": "ABX9TyOJpQcluvkD9ym3mW12oJZf",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/MrEminent42/wordle-ai/blob/main/WordleAI_Test.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## WordleGame"
      ],
      "metadata": {
        "id": "G3KCGZRIuxOj"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "id": "4USCxi7juG70"
      },
      "outputs": [],
      "source": [
        "from enum import IntEnum\n",
        "from termcolor import colored\n",
        "\n",
        "\n",
        "class WordleGame:\n",
        "    def __init__(self, answer):\n",
        "        self.answer = answer.upper()\n",
        "        self.board = []\n",
        "        self.is_over = False\n",
        "        self.win = False\n",
        "\n",
        "    def __repr__(self):\n",
        "        s = \"\"\n",
        "        # in each line\n",
        "        for i, line in enumerate(self.board):\n",
        "            colors = self.get_colors(self.get_line_string(i))\n",
        "            for i, tile in enumerate(line):\n",
        "                # for i, char in enumerate(line):\n",
        "                s += (\n",
        "                    colored(tile.char, \"white\")\n",
        "                    if tile.color == Color.GREY\n",
        "                    else (\n",
        "                        colored(tile.char, \"magenta\")\n",
        "                        if tile.color == Color.YELLOW\n",
        "                        else colored(tile.char, \"green\")\n",
        "                    )\n",
        "                )\n",
        "            s += \"\\n\"\n",
        "\n",
        "        return s\n",
        "\n",
        "    def guess(self, guess):\n",
        "        \"\"\"Takes a five-letter guess, records this guess on the game's board.\n",
        "        Returns the array of Colors with each index corresponding to the color of the letter at that index in the guess\"\"\"\n",
        "        tiles = []\n",
        "        if len(guess) != 5:\n",
        "            raise ValueError(\n",
        "                'Wordle guess must be a 5-letter word. Could not guess with word \"'\n",
        "                + guess\n",
        "                + '\".'\n",
        "            )\n",
        "        # convert everything to upper case\n",
        "        guess = guess.upper()\n",
        "        colors = self.get_colors(guess)\n",
        "        # log guess to board\n",
        "        tiles = [Tile(guess[i], colors[i]) for i in range(5)]\n",
        "        self.board.append(tiles)\n",
        "\n",
        "        print(\"LOGGED COLOR \" + str(int(colors[1])))\n",
        "\n",
        "        # check for game over\n",
        "        if len(self.board) >= 6:\n",
        "            self.is_over = True\n",
        "        elif guess == self.answer:\n",
        "            self.is_over = self.win = True\n",
        "\n",
        "        # give back list of colors\n",
        "        return colors\n",
        "\n",
        "    def get_colors(self, guess):\n",
        "        \"\"\"Takes in a five-letter guess, returns an array of Colors with\n",
        "        each index corresponding to the color of the letter at that index in the guess.\"\"\"\n",
        "        if len(guess) != 5:\n",
        "            raise ValueError(\n",
        "                'Can only find colors for words of length 5. Could not find colors for word \"'\n",
        "                + guess\n",
        "                + '\"'\n",
        "            )\n",
        "        colors = []\n",
        "        # occurrences_left = {char: self.answer.count(char) for char in self.answer}\n",
        "        occurrences_left = {}\n",
        "        # more efficient way of counting num occurences\n",
        "        for char in self.answer:\n",
        "            if char in occurrences_left:\n",
        "                occurrences_left[char] += 1\n",
        "            else:\n",
        "                occurrences_left[char] = 1\n",
        "\n",
        "        for i, char in enumerate(guess):\n",
        "            # if the character is in the correct place, green\n",
        "            if self.answer[i] == char:\n",
        "                colors.append(Color.GREEN)\n",
        "                occurrences_left[char] -= 1\n",
        "            # if the character is in the word, but in the wrong place\n",
        "            elif char in self.answer:\n",
        "                # if there are stil occurences of this char that have not been accounted for\n",
        "                if occurrences_left[char] > 0:\n",
        "                    # append a yellow tile\n",
        "                    colors.append(Color.YELLOW)\n",
        "                    # record that we have accounted for this occurence\n",
        "                    occurrences_left[char] -= 1\n",
        "                else:\n",
        "                    colors.append(Color.GREY)\n",
        "            else:\n",
        "                colors.append(Color.GREY)\n",
        "\n",
        "        return colors\n",
        "\n",
        "    def get_line_string(self, i):\n",
        "        s = \"\"\n",
        "        for tile in self.board[i]:\n",
        "            s += tile.char\n",
        "        return s\n",
        "\n",
        "    def is_over(self):\n",
        "        return self.is_over\n",
        "\n",
        "    def run_game(self):\n",
        "        print(\"Welcome to Wordle-AI!\")\n",
        "        while not self.is_over:\n",
        "            self.guess(input(\"Guess: \"))\n",
        "            print(self)\n",
        "        if self.win:\n",
        "            print(\"Congrats! You found the word in \" + str(len(self.board)) + \" tries.\")\n",
        "        else:\n",
        "            print(\"Darn! You didn't find the word. It was \" + self.answer + \".\")\n",
        "\n",
        "\n",
        "class Color(IntEnum):\n",
        "    GREY = GRAY = 0\n",
        "    YELLOW = 1\n",
        "    GREEN = 2\n",
        "\n",
        "\n",
        "class Tile:\n",
        "    def __init__(self, character, color):\n",
        "        self.char = character\n",
        "        self.color = color"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## DQN\n",
        "Based on https://www.tensorflow.org/agents/tutorials/1_dqn_tutorial"
      ],
      "metadata": {
        "id": "akUNya1tu1KR"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### TF Setup"
      ],
      "metadata": {
        "id": "621fHRwsvWRE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!sudo apt-get update\n",
        "!sudo apt-get install -y xvfb ffmpeg freeglut3-dev\n",
        "!pip install 'imageio==2.4.0'\n",
        "!pip install pyvirtualdisplay\n",
        "!pip install tf-agents[reverb]\n",
        "!pip install pyglet"
      ],
      "metadata": {
        "id": "pgZjknCBuXvG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from __future__ import absolute_import, division, print_function\n",
        "\n",
        "import base64\n",
        "import imageio\n",
        "import IPython\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import PIL.Image\n",
        "import pyvirtualdisplay\n",
        "import reverb\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from tf_agents.agents.dqn import dqn_agent\n",
        "from tf_agents.drivers import py_driver\n",
        "from tf_agents.environments import suite_gym\n",
        "from tf_agents.environments import tf_py_environment\n",
        "from tf_agents.eval import metric_utils\n",
        "from tf_agents.metrics import tf_metrics\n",
        "from tf_agents.networks import sequential\n",
        "from tf_agents.policies import py_tf_eager_policy\n",
        "from tf_agents.policies import random_tf_policy\n",
        "from tf_agents.replay_buffers import reverb_replay_buffer\n",
        "from tf_agents.replay_buffers import reverb_utils\n",
        "from tf_agents.trajectories import trajectory\n",
        "from tf_agents.specs import tensor_spec\n",
        "from tf_agents.utils import common"
      ],
      "metadata": {
        "id": "hiX7OCh0uVdV"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Gym setup"
      ],
      "metadata": {
        "id": "31tTx_F4vYSV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install \"gym>=0.21.0\""
      ],
      "metadata": {
        "id": "NiYkYlGfvbxm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from __future__ import absolute_import\n",
        "from __future__ import division\n",
        "from __future__ import print_function\n",
        "\n",
        "import abc\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "\n",
        "import random\n",
        "\n",
        "from tf_agents.environments import py_environment\n",
        "from tf_agents.environments import tf_environment\n",
        "from tf_agents.environments import tf_py_environment\n",
        "from tf_agents.environments import utils\n",
        "from tf_agents.specs import array_spec\n",
        "from tf_agents.environments import wrappers\n",
        "from tf_agents.environments import suite_gym\n",
        "from tf_agents.trajectories import time_step as ts"
      ],
      "metadata": {
        "id": "xVskUfdZveqt"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Wordle Data Setup"
      ],
      "metadata": {
        "id": "5qx2BXGMmaza"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import requests\n",
        "import io\n",
        "\n",
        "import urllib.request\n",
        "import os\n",
        "from os import path"
      ],
      "metadata": {
        "id": "Asrp22bjzJGV"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "raw_data_source = \"https://raw.githubusercontent.com/tabatkins/wordle-list/main/words\"\n",
        "directory_name = \"WordleData\"\n",
        "file_name = \"wordList.txt\"\n",
        "\n",
        "\n",
        "def retrieve_data():\n",
        "    url_request = urllib.request.urlopen(raw_data_source)\n",
        "    data = url_request.read().decode(\"utf-8\")\n",
        "    return data\n",
        "\n",
        "\n",
        "def store_data():\n",
        "    if os.path.isfile(get_file_path()) == False:\n",
        "        data = retrieve_data()\n",
        "        os.mkdir(directory_name)\n",
        "        file = open(get_file_path(), \"w\")\n",
        "        file.write(data)\n",
        "        file.close\n",
        "\n",
        "\n",
        "def get_file_path():\n",
        "    return path.join(directory_name, file_name)\n"
      ],
      "metadata": {
        "id": "zSCOnRf4mfwS"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "store_data()\n",
        "\n",
        "find_word = {}\n",
        "find_num = {}\n",
        "\n",
        "with open(get_file_path()) as f:\n",
        "  words = f.readlines()\n",
        "\n",
        "for i in range(len(words)):\n",
        "  find_word[i] = words[i]\n",
        "  find_num[words[i]] = i\n",
        "\n",
        "num_words = len(find_word)\n",
        "num_words"
      ],
      "metadata": {
        "id": "lqcyaYfAzLOi",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "94ca901a-0a05-41be-caeb-2b46f768f0af"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "12947"
            ]
          },
          "metadata": {},
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Environment\n",
        "Based on https://www.tensorflow.org/agents/tutorials/2_environments_tutorial#creating_your_own_python_environment"
      ],
      "metadata": {
        "id": "JKfdj7qZvBMJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "wordle_rows = 6 # @param {type:\"integer\"}\n",
        "wordle_cols = 5  # @param {type:\"integer\"}\n",
        "wordle_colors = 3  # @param {type:\"integer\"}\n",
        "\n",
        "english_letters = 26  # @param {type:\"integer\"}"
      ],
      "metadata": {
        "id": "wyA-fOYTpufe"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import random\n",
        "\n",
        "def get_observation_array(game):\n",
        "  obs = np.zeros((wordle_rows * wordle_cols, english_letters, wordle_colors))\n",
        "  # for every tile on the board\n",
        "  for row in range(len(game.board)):\n",
        "    for col in range(len(game.board[row])):\n",
        "      char = game.board[row][col].char\n",
        "      color = game.board[row][col].color\n",
        "      obs[row * 5 + col,ord(char) - 65,int(color)] = 1\n",
        "\n",
        "  return obs\n",
        "\n",
        "def get_reward(game):\n",
        "  sum = 0\n",
        "  for row in game.board:\n",
        "    for tile in row:\n",
        "      sum += int(tile.color) + 1\n",
        "  return sum"
      ],
      "metadata": {
        "id": "FNOPczmKojWq"
      },
      "execution_count": 51,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# test the observation spec\n",
        "game = WordleGame(find_word[random.randint(0,num_words)])\n",
        "game.guess(\"LATER\")\n",
        "print(\"Answer: \", game.answer)\n",
        "obs = get_observation_array(game)\n",
        "print(\"Tile 0 [L]: \", obs[0][ord(\"L\")-65])\n",
        "print(\"Tile 1 [A]: \", obs[1][ord(\"A\")-65])\n",
        "print(\"Tile 2 [T]: \", obs[2][ord(\"T\")-65])\n",
        "print(\"Tile 3 [E]: \", obs[3][ord(\"E\")-65])\n",
        "print(\"Tile 4 [R]: \", obs[4][ord(\"R\")-65])\n",
        "\n",
        "print(\"Current reward:\", get_reward(game))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YmJ2pdkQvYrS",
        "outputId": "d746edfe-38c3-4d13-a45a-088685717e87"
      },
      "execution_count": 56,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "LOGGED COLOR 0\n",
            "Answer:  DOOLY\n",
            "\n",
            "Tile 0 [L]:  [0. 1. 0.]\n",
            "Tile 1 [A]:  [1. 0. 0.]\n",
            "Tile 2 [T]:  [1. 0. 0.]\n",
            "Tile 3 [E]:  [1. 0. 0.]\n",
            "Tile 4 [R]:  [1. 0. 0.]\n",
            "Current reward: 6\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class WordleEnvironment(py_environment.PyEnvironment):\n",
        "  def __init__(self):\n",
        "    self._action_spec = array_spec.BoundedArraySpec(\n",
        "        shape=(), dtype=np.int32, minimum=0, maximum=1, name='action'\n",
        "    )\n",
        "    self._observation_spec = array_spec.BoundedArraySpec(\n",
        "        shape=(wordle_rows * wordle_cols, english_letters, wordle_colors), dtype=np.int32, minimum=0, maximum=1, name='observation'\n",
        "    )\n",
        "    \n",
        "    self.game = WordleGame(find_word[random.randint(0,num_words)])\n",
        "\n",
        "  def action_spec(self):\n",
        "    return self._action_spec\n",
        "  \n",
        "  def observation_spec(self):\n",
        "    return self._observation_spec\n",
        "  \n",
        "  def _reset(self):\n",
        "    self.game = WordleGame(find_word[random.randint(0,num_words)])\n",
        "    return ts.restart(np.array([self._state], dtype=np.int32)) # TODO update\n",
        "\n",
        "  def _step(self, action):\n",
        "    if self.game.is_over: \n",
        "      return self.reset()\n",
        "\n",
        "    # perform the action\n",
        "    colors = self.game.guess(find_word[action])\n",
        "\n",
        "    # return the observation state\n",
        "    # if the game is now over\n",
        "    if self.game.is_over:\n",
        "      return ts.termination(get_observation_array(self.game), get_reward(self.game))\n",
        "\n",
        "\n",
        "    \n",
        "\n",
        "\n",
        "        "
      ],
      "metadata": {
        "id": "RhzLi53QvEIn"
      },
      "execution_count": 62,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "num_iterations = 20000 # @param {type:\"integer\"}\n",
        "\n",
        "initial_collect_steps = 100  # @param {type:\"integer\"}\n",
        "collect_steps_per_iteration =   1# @param {type:\"integer\"}\n",
        "replay_buffer_max_length = 100000  # @param {type:\"integer\"}\n",
        "\n",
        "batch_size = 64  # @param {type:\"integer\"}\n",
        "learning_rate = 1e-3  # @param {type:\"number\"}\n",
        "log_interval = 200  # @param {type:\"integer\"}\n",
        "\n",
        "num_eval_episodes = 10  # @param {type:\"integer\"}\n",
        "eval_interval = 1000  # @param {type:\"integer\"}"
      ],
      "metadata": {
        "id": "V1dlDirUug_O"
      },
      "execution_count": 58,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Training"
      ],
      "metadata": {
        "id": "tJDh45LJx8ys"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "env = WordleEnvironment()\n",
        "print('Observation Spec:')\n",
        "print(env.time_step_spec().observation)\n",
        "print('Reward Spec:')\n",
        "print(env.time_step_spec().reward)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x8Zj-Wdrx_Mv",
        "outputId": "5a709083-c45f-47b3-e667-2a440e769182"
      },
      "execution_count": 64,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Observation Spec:\n",
            "BoundedArraySpec(shape=(30, 26, 3), dtype=dtype('int32'), name='observation', minimum=0, maximum=1)\n",
            "Reward Spec:\n",
            "ArraySpec(shape=(), dtype=dtype('float32'), name='reward')\n"
          ]
        }
      ]
    }
  ]
}