examples/rnn.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A rough copy of https://jaketae.github.io/study/pytorch-rnn/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import warnings\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "from tinygrad.device import Device\n",
    "from tinygrad import Tensor, nn\n",
    "from tinygrad.nn.state import get_parameters\n",
    "from data import fetch_names, name_to_array\n",
    "random.seed(1337)\n",
    "Tensor.manual_seed(1337)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RecurrentNet:\n",
    "    def __init__(self, input_size: int, hidden_size: int, output_size: int):\n",
    "        self.hidden_size = hidden_size\n",
    "        self.in_to_hidden = nn.Linear(input_size + hidden_size, hidden_size)\n",
    "        self.in_to_output = nn.Linear(input_size + hidden_size, output_size)\n",
    "    \n",
    "    def __call__(self, x: Tensor, h: Tensor) -> Tensor:\n",
    "        x = Tensor.cat(x, h, dim=1)\n",
    "        hidden = self.in_to_hidden(x).tanh()\n",
    "        output = self.in_to_output(x).log_softmax()\n",
    "        return output, hidden\n",
    "    \n",
    "    def init_hidden(self):\n",
    "        return Tensor.zeros(1, self.hidden_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set, test_set, char_to_idx, lang_to_label = fetch_names(Tensor, Tensor)\n",
    "num_letters, num_langs = len(char_to_idx), len(lang_to_label)\n",
    "label_to_lang = {label: lang for lang, label in lang_to_label.items()}\n",
    "\n",
    "model = RecurrentNet(num_letters, 256, num_langs)\n",
    "criterion = Tensor.sparse_categorical_crossentropy\n",
    "learning_rate = 0.001"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "loss 0.01: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18063/18063 [10:04<00:00, 29.89it/s]\n",
      "loss 4.60: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 18063/18063 [09:55<00:00, 30.33it/s]\n"
     ]
    }
   ],
   "source": [
    "for _ in range(2):\n",
    "    random.shuffle(train_set)\n",
    "    optimizer = nn.optim.Adam(get_parameters(model), lr=learning_rate)\n",
    "    for i, (name, label) in (t := tqdm(enumerate(train_set), total=len(train_set))):\n",
    "        hidden_state = model.init_hidden()\n",
    "        for char in name:\n",
    "            output, hidden_state = model(char, hidden_state)\n",
    "        loss = criterion(output, label)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        t.set_description(f\"loss {loss.numpy().item():.2f}\")\n",
    "    learning_rate /= 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test set accuracy is 0.0\n"
     ]
    }
   ],
   "source": [
    "num_correct = 0\n",
    "for name, label in test_set:\n",
    "    hidden_state = model.init_hidden()\n",
    "    for char in name:\n",
    "        output, hidden_state = model(char, hidden_state)\n",
    "    pred = output.argmax(-1).numpy()\n",
    "    num_correct += (pred == label).sum()\n",
    "print(f\"test set accuracy is {num_correct/len(test_set)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict(name: str) -> str:\n",
    "    tensor_name = Tensor(name_to_array(name, char_to_idx))\n",
    "    hidden_state = model.init_hidden()\n",
    "    for char in tensor_name:\n",
    "        output, hidden_state = model(char, hidden_state)\n",
    "    pred = output.numpy().argmax()\n",
    "    return label_to_lang[pred]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jake: Russian\n",
      "Qin: Chinese\n",
      "Fernando: Italian\n",
      "Demirkan: Russian\n"
     ]
    }
   ],
   "source": [
    "for name in (\"Jake\", \"Qin\", \"Fernando\", \"Demirkan\"):\n",
    "    print(f\"{name}: {predict(name)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "7e1998ff7f8aa20ada591c520b972326324e5ea05489af9e422744c7c09f6dad"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}