CodeBOT.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Code BOT \n",
        "\n",
        "> by Avishake Adhikary"
      ],
      "metadata": {
        "id": "H6yxPZRUPMQI"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Installing Dependencies"
      ],
      "metadata": {
        "id": "7q73qlDMPMVJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install transformers "
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LixmAIFvPSv3",
        "outputId": "a8b24117-3e2d-46eb-c18d-63712dab6b9f"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (4.29.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.12.0)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.14.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.14.1)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (1.22.4)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (23.1)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (6.0)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2022.10.31)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers) (2.27.1)\n",
            "Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.13.3)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers) (4.65.0)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.14.1->transformers) (2023.4.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.14.1->transformers) (4.5.0)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (1.26.15)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2022.12.7)\n",
            "Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2.0.12)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.4)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Import Dependencies"
      ],
      "metadata": {
        "id": "AfszAZ2KPlyv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
        "import torch"
      ],
      "metadata": {
        "id": "iOBqV9oBPocr"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Initialize Variables"
      ],
      "metadata": {
        "id": "7QVdXnLDPsDC"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "checkpoint = \"Salesforce/codegen2-1B\"\n",
        "\n",
        "# checkpoint = \"Salesforce/codegen2-7B\" #uncomment in case of other model\n",
        "# checkpoint = \"Salesforce/codegen2-3_7B\" #uncomment in case of other model\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
      ],
      "metadata": {
        "id": "u-aqhqtNP3to"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Initialize Tokenizer and Model"
      ],
      "metadata": {
        "id": "fAD_wnsbQPC4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n",
        "model = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True, revision=\"main\")"
      ],
      "metadata": {
        "id": "wXhifrGsQRSF"
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Create Prompt"
      ],
      "metadata": {
        "id": "6CUgGWo7Q4Sk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "prompt = str(input(\"Enter your prompt: \"))\n",
        "input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids\n",
        "\n",
        "generated_ids = model.generate(input_ids, max_length=1024) # Tweak the maximum length of generated id's to get better code.\n",
        "\n",
        "print(tokenizer.decode(generated_ids[0], skip_special_tokens=True))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-rnSyV3nQ3pu",
        "outputId": "66b1f085-83da-4de2-97f3-8156bd89d5bc"
      },
      "execution_count": 6,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Enter your prompt: def hello_world():\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
            "Setting `pad_token_id` to `eos_token_id`:2 for open-end generation.\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "def hello_world():\n",
            "    print(\"Hello World!\")\n",
            "\n",
            "\n",
            "if __name__ == \"__main__\":\n",
            "    hello_world()\n",
            "<|python|>#\n"
          ]
        }
      ]
    }
  ]
}