{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": " gradio_with_CodeGen.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Interacting with [CodeGen](https://github.com/salesforce/CodeGen/)\n",
        "\n",
        " "
      ],
      "metadata": {
        "id": "UFoSzgxnvrlb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        " "
      ],
      "metadata": {
        "id": "5VtHpiGpqLWF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wPwrvvMlkTVA"
      },
      "outputs": [],
      "source": [
        " \n",
        "!git clone https://github.com/salesforce/CodeGen\n",
        "%cd CodeGen\n",
        "!pip install --upgrade pip setuptools\n",
        "!pip install gradio\n",
        "!pip install -r requirements.txt"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        " \n",
        "chosen_model = \"codegen-350M-nl\" #@param [\"codegen-350M-nl\", \"codegen-350M-multi\", \"codegen-350M-mono\", \"codegen-2B-nl\", \"codegen-2B-multi\", \"codegen-2B-mono\", \"codegen-6B-nl\", \"codegen-6B-multi\", \"codegen-6B-mono\", \"codegen-16B-nl\", \"codegen-16B-multi\", \"codegen-16B-mono\"]\n",
        "fp16 = True #param {type:\"boolean\"}\n",
        "\n",
        "import os\n",
        "\n",
        "if not os.path.exists(f'./checkpoints/{chosen_model}'):\n",
        "  !wget -P checkpoints https://storage.googleapis.com/sfr-codegen-research/checkpoints/{chosen_model}.tar.gz && tar -xvf checkpoints/{chosen_model}.tar.gz -C checkpoints/\n",
        "\n",
        "\n",
        "import torch\n",
        "from jaxformer.hf.sample import truncate as do_truncate\n",
        "from jaxformer.hf.sample import set_env, set_seed, print_time, create_model, create_custom_gpt2_tokenizer, create_tokenizer, sample\n",
        "\n",
        "# (0) constants\n",
        "\n",
        "models_nl = ['codegen-350M-nl', 'codegen-2B-nl', 'codegen-6B-nl', 'codegen-16B-nl']\n",
        "models_pl = ['codegen-350M-multi', 'codegen-2B-multi', 'codegen-6B-multi', 'codegen-16B-multi', 'codegen-350M-mono', 'codegen-2B-mono', 'codegen-6B-mono', 'codegen-16B-mono']\n",
        "models = models_nl + models_pl\n",
        "\n",
        "\n",
        "# (2) preamble\n",
        "\n",
        "set_env()\n",
        "\n",
        "pad = 50256\n",
        "# device = torch.device('cuda:0')\n",
        "device = torch.device(\"cpu\")\n",
        "ckpt = f'./checkpoints/{chosen_model}'\n",
        "\n",
        "# if device.type == \"cpu\":\n",
        "#   print()\n",
        "#   print(\"force full precision for cpu!!\")\n",
        "#   print()\n",
        "  \n",
        "fp16 = False\n",
        "\n",
        "\n",
        "# (3) load\n",
        "\n",
        "with print_time('loading parameters'):\n",
        "  model = create_model(ckpt=ckpt, fp16=fp16).to(device)\n",
        "\n",
        "\n",
        "with print_time('loading tokenizer'):\n",
        "  if chosen_model in models_pl:\n",
        "    tokenizer = create_custom_gpt2_tokenizer()\n",
        "  else:\n",
        "    tokenizer = create_tokenizer()\n",
        "  tokenizer.padding_side = 'left'\n",
        "  tokenizer.pad_token = pad"
      ],
      "metadata": {
        "id": "xZbWFH1hkf9i"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def codegen(context):\n",
        " #param {type:\"string\"}\n",
        "  \n",
        "  rng_seed = 42 #param {type:\"integer\"}\n",
        "  rng_deterministic = True #param {type:\"boolean\"}\n",
        "  p = 0.95 #param {type:\"number\"}\n",
        "  t = 0.1 #param {type:\"number\"}\n",
        "  max_length = 128 #param {type:\"integer\"}\n",
        "  batch_size = 1 #param {type:\"integer\"}\n",
        "  set_seed(rng_seed, deterministic=rng_deterministic)\n",
        "\n",
        "  # (4) sample\n",
        "\n",
        "  with print_time('sampling'):\n",
        "    completion = sample(device=device, model=model, tokenizer=tokenizer, context=context, pad_token_id=pad, num_return_sequences=batch_size, temp=t, top_p=p, max_length_sample=max_length)[0]\n",
        "    truncation = do_truncate(completion)\n",
        "\n",
        "    # print('=' * 100)\n",
        "    # print(completion)\n",
        "    # print('=' * 100)\n",
        "    # print(context+truncation)\n",
        "    # print('=' * 100)\n",
        "      \n",
        "\n",
        "  return completion\n",
        "\n",
        "# !python -m jaxformer.hf.sample --model $chosen_model \\\n",
        "#                  --rng-seed $rng_seed \\\n",
        "#                  --p $p \\\n",
        "#                  --t $t \\\n",
        "#                  --max-length $max_length \\\n",
        "#                  --batch-size $batch_size \\\n",
        "#                  --context '$context'\n",
        "    "
      ],
      "metadata": {
        "id": "YN2xY4xmkss0"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "# context = \"def hello_world():\"\n",
        "# codegen(context)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 87
        },
        "id": "HcaxDQNRNB8u",
        "outputId": "c28a1c68-1e03-4550-8658-f749350ab52d"
      },
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "sampling\n",
            "sampling took 17.08s\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'\\n        print(\"Hello world\")\\n        hello_world()\\n        hello_world()\\n        hello_world()\\n        hello_world()\\n        hello_world()\\n        hello_world()\\n        hello_world()\\n        hello_world()\\n        hello_world()\\n      '"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "import numpy as np\n",
        "import gradio as gr\n",
        " \n",
        "\n",
        "iface = gr.Interface(\n",
        "    codegen,\n",
        "    [   gr.inputs.Textbox(type='str',   label=\"input prompt\"),\n",
        "    ],\n",
        "    \"text\",\n",
        ")\n",
        "\n",
        "iface.launch(debug=True)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QLvCESs0NmZC",
        "outputId": "9ac652a2-2f80-402a-a079-e56fadce9d29"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Keyboard interruption in main thread... closing server.\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(<fastapi.applications.FastAPI at 0x7fbdbd91c090>,\n",
              " 'http://127.0.0.1:7860/',\n",
              " 'https://33801.gradio.app')"
            ]
          },
          "metadata": {},
          "execution_count": 15
        }
      ]
    }
  ]
}