ChatRWKV.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ymdysk/chatrwkv-notebook/blob/main/ChatRWKV.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hoLtM9B2vKwm"
      },
      "source": [
        "# Notebook for running ChatRWKV with Google Colab\n",
        "## How it works\n",
        "\n",
        "- Mount Google Drive and clone the ChatRWKV Git repository in the src folder.\n",
        "- Go to the ChatRWKV/v2/ folder and download a model.\n",
        "- Set variables related to the model in the form, replace some contents of chat.py and save it in chat-notebook.py to load the model.\n",
        "- Set variables related to text generation and run the chat.\n",
        "\n",
        "## Tips\n",
        "- If you want to use GPU, go to Google Colab's \"Runtime\" menu -> \"Change Runtime Type\" -> \"Hardware Accelerator\" and select \"GPU\" and save the setting.\n",
        "- If GPU is not available, setting strategy = 'cpu fp32', RWKV_CUDA_ON = 0 will work, but the generation speed will be slower.\n",
        "- The model to be used can be either a 1B parameter model of about 3GB or a 3B parameter model of about 6GB by selecting MODEL_URL or entering an arbitrary URL. Using a model with a higher number of parameters will result in a more intelligent response, but will also require a larger storage size and higher spec runtime for VRAM, RAM, etc.\n",
        "- The free version of Google Colaboratory seems to be able to run 1B and 3B models as is.\n",
        "- Google's free storage is 15GB, but it is easy to get tight if you store models in Google Drive.\n",
        "\n",
        "## License\n",
        "- This notebook  \n",
        "  Copyright 2023 Yosuke Yamada  \n",
        "  Licensed under the Apache License, Version 2.0  \n",
        "  http://www.apache.org/licenses/LICENSE-2.0\n",
        "- Please check the license of software and models used/downloaded from the notebook individually.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Check the environment"
      ],
      "metadata": {
        "id": "DmPHwlbpGBGZ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jvf3g-UOHp3D",
        "outputId": "d540b96a-a359-4a37-a48e-282ae61bfb85"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Sat Apr 22 16:38:49 2023       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   47C    P8    10W /  70W |      0MiB / 15360MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "# Check the status of the CUDA environment on NVIDIA's system management interface (for GPUs, not necessary if only CPU is used)\n",
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kHuA67Z8I-sw",
        "outputId": "2bae21d7-300e-4839-8f44-ed78a7a3695e"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "nvcc: NVIDIA (R) Cuda compiler driver\n",
            "Copyright (c) 2005-2022 NVIDIA Corporation\n",
            "Built on Wed_Sep_21_10:33:58_PDT_2022\n",
            "Cuda compilation tools, release 11.8, V11.8.89\n",
            "Build cuda_11.8.r11.8/compiler.31833905_0\n"
          ]
        }
      ],
      "source": [
        "# Check the version of the Cuda compiler (for GPU, not necessary if using CPU only).\n",
        "!nvcc --version"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Environment Settings"
      ],
      "metadata": {
        "id": "oYMQDqw4GRYD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# If you use the default torch in the Google Colab environment, an error may occur if you set RWKV_CUDA_ON = 1 to speed up the process, so re-install the torch (first time only).\n",
        "!pip uninstall -y torch\n",
        "!pip install torch==2.0.0+cu118 -f https://download.pytorch.org/whl/torch_stable.html"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jr70F9lc_r-l",
        "outputId": "2d73b3cb-b38f-4e3f-a1d8-8de6a150309c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Found existing installation: torch 2.0.0+cu118\n",
            "Uninstalling torch-2.0.0+cu118:\n",
            "  Successfully uninstalled torch-2.0.0+cu118\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Looking in links: https://download.pytorch.org/whl/torch_stable.html\n",
            "Collecting torch==2.0.0+cu118\n",
            "  Downloading https://download.pytorch.org/whl/cu118/torch-2.0.0%2Bcu118-cp39-cp39-linux_x86_64.whl (2267.3 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.3/2.3 GB\u001b[0m \u001b[31m772.7 kB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.9/dist-packages (from torch==2.0.0+cu118) (3.11.0)\n",
            "Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.9/dist-packages (from torch==2.0.0+cu118) (2.0.0)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.9/dist-packages (from torch==2.0.0+cu118) (4.5.0)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.9/dist-packages (from torch==2.0.0+cu118) (3.1.2)\n",
            "Requirement already satisfied: networkx in /usr/local/lib/python3.9/dist-packages (from torch==2.0.0+cu118) (3.1)\n",
            "Requirement already satisfied: sympy in /usr/local/lib/python3.9/dist-packages (from torch==2.0.0+cu118) (1.11.1)\n",
            "Requirement already satisfied: cmake in /usr/local/lib/python3.9/dist-packages (from triton==2.0.0->torch==2.0.0+cu118) (3.25.2)\n",
            "Requirement already satisfied: lit in /usr/local/lib/python3.9/dist-packages (from triton==2.0.0->torch==2.0.0+cu118) (16.0.1)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.9/dist-packages (from jinja2->torch==2.0.0+cu118) (2.1.2)\n",
            "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.9/dist-packages (from sympy->torch==2.0.0+cu118) (1.3.0)\n",
            "Installing collected packages: torch\n",
            "Successfully installed torch-2.0.0+cu118\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6Bb6gKeVT0Pe",
        "outputId": "5fbe91e2-1027-4625-cc0a-011af9c22fad"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "True\n",
            "0\n"
          ]
        }
      ],
      "source": [
        "# Check if CUDA is available in Pytorch. If so, \"True\" and the device number will be returned. If it fails, run the cell above again. If only CPU-only (no GPU) runtime is used, there is no need to run this cell.\n",
        "import torch\n",
        "print(torch.cuda.is_available())\n",
        "print(torch.cuda.current_device())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1UYqFX-DvU4C",
        "outputId": "134dc0eb-01b5-49a6-b472-2d15d8be30d8"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount(\"/content/drive\", force_remount=True).\n",
            "/content/drive/My Drive/src\n"
          ]
        }
      ],
      "source": [
        "# Mount Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "# If there is no source code folder, create it and cd it.\n",
        "import os\n",
        "os.makedirs(\"/content/drive/My Drive/src\", exist_ok=True)\n",
        "%cd '/content/drive/My Drive/src'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LSMAkkVbG4K0",
        "outputId": "3590074a-6166-483a-b604-eaa66b597e1f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'ChatRWKV'...\n",
            "remote: Enumerating objects: 1340, done.\u001b[K\n",
            "remote: Counting objects: 100% (216/216), done.\u001b[K\n",
            "remote: Compressing objects: 100% (149/149), done.\u001b[K\n",
            "remote: Total 1340 (delta 88), reused 162 (delta 52), pack-reused 1124\u001b[K\n",
            "Receiving objects: 100% (1340/1340), 26.97 MiB | 17.77 MiB/s, done.\n",
            "Resolving deltas: 100% (723/723), done.\n"
          ]
        }
      ],
      "source": [
        "# Get source code of ChatRWKV by git clone (first time only)\n",
        "!git clone https://github.com/BlinkDL/ChatRWKV"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3KrxdGeUGJ7X",
        "outputId": "f8e9128a-4b1f-4d28-edfc-de9d52ebf692"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "/content/drive/My Drive/src/ChatRWKV/v2\n"
          ]
        }
      ],
      "source": [
        "# Move to ChatRWKV/v2 folder\n",
        "%cd 'ChatRWKV/v2'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZG4WhkbLJdpy",
        "outputId": "06994e0f-33df-4c1a-a52a-b09c3f204772"
      },
      "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: rwkv in /usr/local/lib/python3.9/dist-packages (0.7.3)\n",
            "Requirement already satisfied: ninja in /usr/local/lib/python3.9/dist-packages (1.11.1)\n",
            "Requirement already satisfied: tokenizers>=0.13.2 in /usr/local/lib/python3.9/dist-packages (from rwkv) (0.13.3)\n"
          ]
        }
      ],
      "source": [
        "# Install rwkv and ninja package by pip\n",
        "!pip install rwkv ninja"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Prepare Model"
      ],
      "metadata": {
        "id": "bXvywPMzGlXQ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8hTTbe9fB0Cn",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "238af518-c3fa-464e-8dad-b9d5bf66a518"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "** Resuming transfer from byte position 3030279730\n",
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "\r  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0\r100  1217  100  1217    0     0   4791      0 --:--:-- --:--:-- --:--:--  4791\n",
            "\r100    49  100    49    0     0    173      0 --:--:-- --:--:-- --:--:--   173\n"
          ]
        }
      ],
      "source": [
        "# Specify model and download\n",
        "# Raven  https://huggingface.co/BlinkDL/rwkv-4-raven/tree/main\n",
        "# Others https://huggingface.co/BlinkDL\n",
        "MODEL_URL = 'https://huggingface.co/BlinkDL/rwkv-4-raven/resolve/main/RWKV-4-Raven-1B5-v10-Eng99%25-Other1%25-20230418-ctx4096.pth' #@param ['https://huggingface.co/BlinkDL/rwkv-4-raven/resolve/main/RWKV-4-Raven-1B5-v10-Eng99%25-Other1%25-20230418-ctx4096.pth', 'https://huggingface.co/BlinkDL/rwkv-4-raven/resolve/main/RWKV-4-Raven-3B-v8-EngAndMore-20230408-ctx4096.pth\"]  {allow-input: true}\n",
        "!curl -OLC - $MODEL_URL\n",
        "# MODEL_NAME is the file name after the \"/\" sign near the end of MODEL_URL, with the \".pth\" extension removed\n",
        "MODEL_NAME = MODEL_URL[MODEL_URL.rfind('/') + 1:].rstrip('.pth')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gm3QS13qlax2"
      },
      "outputs": [],
      "source": [
        "# Settings for RWKV model behavior\n",
        "strategy = 'cuda fp16' #@param ['cpu fp32', 'cuda fp16', 'cuda:0 fp16 -> cuda:1 fp16', 'cuda fp16i8 *10 -> cuda fp16', 'cuda fp16i8', 'cuda fp16i8 -> cpu fp32 *10', 'cuda fp16i8 *10+'] {allow-input: true}\n",
        "# 'cpu fp32' # If CUDA is not available and CPU is used\n",
        "# 'cuda fp16' # Default value when CUDA is available\n",
        "# 'cuda:0 fp16 -> cuda:1 fp16' # If two GPUs can be used\n",
        "# 'cuda fp16i8 *10 -> cuda fp16' # first 10 layers cuda int8 quantization, rest cuda fp16i8\n",
        "# 'cuda fp16i8' # all layers cuda int8 quantization\n",
        "# 'cuda fp16i8 -> cpu fp32 *10' # first is cuda fp16i8, subsequent 10 layers are cpu fp32\n",
        "# 'cuda fp16i8 *10+' # first 10 layers cuda int8 quantize, rest dynamically load as needed\n",
        "\n",
        "# 1 if CUDA is available, 0 if not\n",
        "RWKV_CUDA_ON = \"1\" #@param [0, 1]\n",
        "\n",
        "# Language used in chat\n",
        "CHAT_LANG = 'English' #@param [\"Japanese\", \"English\", \"Chinese\"]\n",
        "# English\n",
        "# Chinese\n",
        "# Japanese\n",
        "\n",
        "# Context length of model\n",
        "ctx_len = 1024 #@param {type:\"integer\"}\n",
        "\n",
        "# Create dictionary to replace contents of chat.py\n",
        "replacements = {\n",
        "    'args.strategy = .*' : 'args.strategy = \\'' + strategy + '\\'',\n",
        "    'os\\.environ\\[\\\"RWKV_CUDA_ON\\\"\\] = \\'.' : 'os.environ[\"RWKV_CUDA_ON\"] = \\'' + str(RWKV_CUDA_ON),\n",
        "    'CHAT_LANG = .*' : 'CHAT_LANG = \\'' + CHAT_LANG + '\\'',\n",
        "    'args.MODEL_NAME = .*' : 'args.MODEL_NAME = \\'' + MODEL_NAME + '\\'',\n",
        "    'args.ctx_len = .*' : 'args.ctx_len = \\'' + str(ctx_len) + '\\'',\n",
        "    'current_path = os\\.path\\.dirname\\(os\\.path\\.abspath\\(__file__\\)\\)' : 'current_path = os.getcwd()',\n",
        "    'while True:\\s+msg = prompt.+\\s+if len\\(msg.+\\s+on_message\\(msg\\)\\s+else:\\s+print\\(.+' : ''\n",
        "}\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FkBn_XeEJinG"
      },
      "outputs": [],
      "source": [
        "# Replace variables in chat.py and save in chat-notebook.py\n",
        "import re\n",
        "\n",
        "with open('chat.py', encoding='utf-8') as f:\n",
        "    text = f.read()\n",
        "\n",
        "for old, new in replacements.items():\n",
        "    pattern = re.compile(r'^(\\s*)' + old, flags=re.MULTILINE)\n",
        "    text = pattern.sub(r'\\1' + new, text)\n",
        "\n",
        "with open('chat-notebook.py', 'w', encoding='utf-8') as f:\n",
        "    f.write(text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "l1PB9PyA28k-",
        "outputId": "183748f0-764f-45ab-fac1-cb6d2de908b0"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "\n",
            "ChatRWKV v2 https://github.com/BlinkDL/ChatRWKV\n",
            "\n",
            "English - cuda fp16 - /content/drive/MyDrive/src/ChatRWKV/v2/prompt/default/English-2.py\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Using /root/.cache/torch_extensions/py39_cu118 as PyTorch extensions root...\n",
            "Detected CUDA files, patching ldflags\n",
            "Emitting ninja build file /root/.cache/torch_extensions/py39_cu118/wkv_cuda/build.ninja...\n",
            "Building extension module wkv_cuda...\n",
            "Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)\n",
            "Loading extension module wkv_cuda...\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loading model - RWKV-4-Raven-1B5-v10-Eng99%25-Other1%25-20230418-ctx4096\n",
            "RWKV_JIT_ON 1 RWKV_CUDA_ON 1 RESCALE_LAYER 6\n",
            "\n",
            "Loading RWKV-4-Raven-1B5-v10-Eng99%25-Other1%25-20230418-ctx4096.pth ...\n",
            "Strategy: (total 24+1=25 layers)\n",
            "* cuda [float16, float16], store 25 layers\n",
            "0-cuda-float16-float16 1-cuda-float16-float16 2-cuda-float16-float16 3-cuda-float16-float16 4-cuda-float16-float16 5-cuda-float16-float16 6-cuda-float16-float16 7-cuda-float16-float16 8-cuda-float16-float16 9-cuda-float16-float16 10-cuda-float16-float16 11-cuda-float16-float16 12-cuda-float16-float16 13-cuda-float16-float16 14-cuda-float16-float16 15-cuda-float16-float16 16-cuda-float16-float16 17-cuda-float16-float16 18-cuda-float16-float16 19-cuda-float16-float16 20-cuda-float16-float16 21-cuda-float16-float16 22-cuda-float16-float16 23-cuda-float16-float16 24-cuda-float16-float16 \n",
            "emb.weight                        f16      cpu  50277  2048 \n",
            "blocks.0.ln1.weight               f16   cuda:0   2048       \n",
            "blocks.0.ln1.bias                 f16   cuda:0   2048       \n",
            "blocks.0.ln2.weight               f16   cuda:0   2048       \n",
            "blocks.0.ln2.bias                 f16   cuda:0   2048       \n",
            "blocks.0.att.time_decay           f32   cuda:0   2048       \n",
            "blocks.0.att.time_first           f32   cuda:0   2048       \n",
            "blocks.0.att.time_mix_k           f16   cuda:0   2048       \n",
            "blocks.0.att.time_mix_v           f16   cuda:0   2048       \n",
            "blocks.0.att.time_mix_r           f16   cuda:0   2048       \n",
            "blocks.0.att.key.weight           f16   cuda:0   2048  2048 \n",
            "blocks.0.att.value.weight         f16   cuda:0   2048  2048 \n",
            "blocks.0.att.receptance.weight    f16   cuda:0   2048  2048 \n",
            "blocks.0.att.output.weight        f16   cuda:0   2048  2048 \n",
            "blocks.0.ffn.time_mix_k           f16   cuda:0   2048       \n",
            "blocks.0.ffn.time_mix_r           f16   cuda:0   2048       \n",
            "blocks.0.ffn.key.weight           f16   cuda:0   2048  8192 \n",
            "blocks.0.ffn.receptance.weight    f16   cuda:0   2048  2048 \n",
            "blocks.0.ffn.value.weight         f16   cuda:0   8192  2048 \n",
            "............................................................................................................................................................................................................................................................................................................................................................................................................\n",
            "blocks.23.ln1.weight              f16   cuda:0   2048       \n",
            "blocks.23.ln1.bias                f16   cuda:0   2048       \n",
            "blocks.23.ln2.weight              f16   cuda:0   2048       \n",
            "blocks.23.ln2.bias                f16   cuda:0   2048       \n",
            "blocks.23.att.time_decay          f32   cuda:0   2048       \n",
            "blocks.23.att.time_first          f32   cuda:0   2048       \n",
            "blocks.23.att.time_mix_k          f16   cuda:0   2048       \n",
            "blocks.23.att.time_mix_v          f16   cuda:0   2048       \n",
            "blocks.23.att.time_mix_r          f16   cuda:0   2048       \n",
            "blocks.23.att.key.weight          f16   cuda:0   2048  2048 \n",
            "blocks.23.att.value.weight        f16   cuda:0   2048  2048 \n",
            "blocks.23.att.receptance.weight   f16   cuda:0   2048  2048 \n",
            "blocks.23.att.output.weight       f16   cuda:0   2048  2048 \n",
            "blocks.23.ffn.time_mix_k          f16   cuda:0   2048       \n",
            "blocks.23.ffn.time_mix_r          f16   cuda:0   2048       \n",
            "blocks.23.ffn.key.weight          f16   cuda:0   2048  8192 \n",
            "blocks.23.ffn.receptance.weight   f16   cuda:0   2048  2048 \n",
            "blocks.23.ffn.value.weight        f16   cuda:0   8192  2048 \n",
            "ln_out.weight                     f16   cuda:0   2048       \n",
            "ln_out.bias                       f16   cuda:0   2048       \n",
            "head.weight                       f16   cuda:0   2048 50277 \n",
            "\n",
            "Run prompt...\n",
            "Commands:\n",
            "say something --> chat with bot. use \\n for new line.\n",
            "+ --> alternate chat reply\n",
            "+reset --> reset chat\n",
            "\n",
            "+gen YOUR PROMPT --> free single-round generation with any prompt. use \\n for new line.\n",
            "+i YOUR INSTRUCT --> free single-round generation with any instruct. use \\n for new line.\n",
            "+++ --> continue last free generation (only for +gen / +i)\n",
            "++ --> retry last free generation (only for +gen / +i)\n",
            "\n",
            "Now talk with the bot and enjoy. Remember to +reset periodically to clean up the bot's memory. Use RWKV-4 14B (especially https://huggingface.co/BlinkDL/rwkv-4-raven) for best results.\n",
            "\n",
            "English - RWKV-4-Raven-1B5-v10-Eng99%25-Other1%25-20230418-ctx4096 - cuda fp16\n",
            "\n",
            "The following is a coherent verbose detailed conversation between a girl named Alice and her friend Bob. Alice is very intelligent, creative and friendly. Alice is unlikely to disagree with Bob, and Alice doesn't like to ask Bob questions. Alice likes to tell Bob a lot about herself and her opinions. Alice usually gives Bob kind, helpful and informative advices.\n",
            "\n",
            "Bob: Hello Alice, how are you doing?\n",
            "Alice: Hi! Thanks, I'm fine. What about you?\n",
            "\n",
            "Bob: I am fine. It's nice to see you. Look, here is a store selling tea and juice.\n",
            "Alice: Sure. Let's go inside. I would like to have some Mocha latte, which is my favourite!\n",
            "\n",
            "Bob: What is it?\n",
            "Alice: Mocha latte is usually made with espresso, milk, chocolate, and frothed milk. Its flavors are frequently sweet.\n",
            "\n",
            "Bob: Sounds tasty. I'll try it next time. Would you like to chat with me for a while?\n",
            "Alice: Of course! I'm glad to answer your questions or give helpful advices. You know, I am confident with my expertise. So please go ahead!\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Load the model\n",
        "execfile(\"chat-notebook.py\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Chat"
      ],
      "metadata": {
        "id": "SW7tRuFSGxj4"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WCcLpH_GpXEI"
      },
      "outputs": [],
      "source": [
        "# Settings related to sentence generation\n",
        "\n",
        "# Short response length for chat\n",
        "CHAT_LEN_SHORT = 40 #@param {type:\"integer\"}\n",
        "# Long response length for chat\n",
        "CHAT_LEN_LONG = 150 #@param {type:\"integer\"}\n",
        "# length of freely generated sentences\n",
        "FREE_GEN_LEN = 256 #@param {type:\"integer\"}\n",
        "\n",
        "# For better chat & QA quality: reduce temp, reduce top-p, increase repetition penalties\n",
        "# Explanation: https://platform.openai.com/docs/api-reference/parameter-details\n",
        "\n",
        "# GEN_TEMP and GEN_TOP_P: smaller values increase accuracy, larger values increase diversity\n",
        "GEN_TEMP = 1.1 #@param {type:\"number\"} # sometimes it's a good idea to increase temp. try it\n",
        "GEN_TOP_P = 0.7 #@param {type:\"number\"}\n",
        "# GEN_alpha_presence, GEN_alpha_frequency: Penalty for presence and frequency of repeated strings. Larger values suppress repetition.\n",
        "GEN_alpha_presence = 0.2 #@param {type:\"number\"} # Presence Penalty\n",
        "GEN_alpha_frequency = 0.2 #@param {type:\"number\"} # Frequency Penalty\n",
        "# AVOID_REPEAT: character to prevent repetition\n",
        "AVOID_REPEAT = '，：？！' #@param {type:\"string\"}\n",
        "# Chunk length to split input\n",
        "CHUNK_LEN = 256 #@param {type:\"integer\"} # split input into chunks to save VRAM (shorter -> slower)\n",
        "\n",
        "AVOID_REPEAT_TOKENS = []\n",
        "for i in AVOID_REPEAT:\n",
        "    dd = pipeline.encode(i)\n",
        "    assert len(dd) == 1\n",
        "    AVOID_REPEAT_TOKENS += dd\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oFnPmwyp5fCA",
        "outputId": "9eca10c4-f829-4d08-a7cf-f211a3ee8950"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Bob: Hi. Please introduce yourself.\n",
            "Alice: Hello, my name is Alice and I am a strong believer in logic and rationality. I love the beauty of the world and all its complexity. I am interested in politics, economics, and science. I have an excellent understanding of philosophy and I am an active member of a few philosophical societies.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Running a chat\n",
        "while True:\n",
        "    msg = input(f'{user}{interface} ')\n",
        "    if len(msg.strip()) > 0:\n",
        "        on_message(msg)\n",
        "    else:\n",
        "        print('Error: please say something')"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "toc_visible": true,
      "provenance": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}