experiments/nlp_experiments/3.4 Chat with Mamba (base).ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "748447e6-88b1-4388-8f33-17ba04755e3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.insert(0, '../../')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c37eb38b-7d74-46a7-b6ab-f23752e46dec",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from typing import List, Tuple, Union\n",
    "\n",
    "import candle\n",
    "import experiments.textgenutils as gutils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4ce2669c-7fb1-47bf-b808-5f0ff4393b22",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fea36b8-ddce-4525-9ac8-e8dd5581f121",
   "metadata": {},
   "source": [
    "## (1) Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9194cfc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['TRANSFORMERS_CACHE'] = '/mnt/disks/disk1'\n",
    "MODEL_ID = 'state-spaces/mamba-2.8b-slimpj'\n",
    "\n",
    "model = candle.models.mamba.Mamba.from_pretrained(MODEL_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8777b772-6491-4c13-8e54-d9fd784c5d3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b')  # Get HF tokenizer. I'm a little too lazy to reimplement yet another tokenizer :-)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d32362b7",
   "metadata": {},
   "source": [
    "## (2) Have a conversation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6d235056-a55d-42b5-8055-ac817fd60f3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "USER_NAME = 'John'\n",
    "ASST_NAME = 'Mamba'\n",
    "\n",
    "SYSTEM_MESSAGE = (f'Two friends, {USER_NAME} and {ASST_NAME}, are having a conversation. '\n",
    "                  f'{ASST_NAME} is friendly, talkative, and loves to ask {USER_NAME} questions.')\n",
    "chat_template = candle.nlp.chattemplates.SimpleConversationTemplate(USER_NAME, ASST_NAME, SYSTEM_MESSAGE)  \n",
    "model.block_size = np.inf  # To maintain compability with generation scripts\n",
    "\n",
    "def start_conversation(profile_pic: str = '🙂'):\n",
    "    gutils.interactive_conversation(\n",
    "        model, chat_template, tokenizer,\n",
    "        user_name=USER_NAME,\n",
    "        profile_pic=profile_pic,\n",
    "        user_bg_color='yellow',\n",
    "        asst_name=ASST_NAME,\n",
    "        asst_profile_pic='🐍',\n",
    "        asst_bg_color='bluegreen',\n",
    "        max_response_length=512,\n",
    "        top_k=100,\n",
    "        top_p=0.90,\n",
    "        temperature=1.2,\n",
    "        stop_token_idx=tokenizer.eos_token_id,\n",
    "        stop_strings={'\\n': 1, '\\.|\\!|\\?': 3},\n",
    "        use_kv_cache=False\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "48275cb9-e072-49a3-83c7-3086c4a43839",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m< You are now talking with Mamba. Send 'bye' to exit, 'clear' to reset cache. >\u001b[0m\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m Good evening!\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m Hi, how are you?\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m I'm good! What are your favourite things to do on a pleasant evening such as this one?\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m Oh, me? Oh, I think I'll run by the bakery to buy bread to go with my supper..\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m Nice! What are you eating for supper?\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m Eh, it's hamburger steaks. The house special.\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m Sounds yummy.\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m And I also like a green salad with a cheese, cucumber, and bacon vinaigrette. And for dessert I like apple pie with vanilla ice cream. A vanilla ice cream sundae on the side\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m Bye\n",
      "\u001b[1m\n",
      "< / end of conversation. >\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "start_conversation(profile_pic='👦🏻')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "353799ee-5a2e-457e-a957-dbb8cb7b3ba8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m< You are now talking with Mamba. Send 'bye' to exit, 'clear' to reset cache. >\u001b[0m\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m Hi Mamba, what's the best laptop to buy in 2023?\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m The Acer.\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m Why?\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m Because it's better than the Asus and Samsung. And they are the best too.\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m What specs should I look for in a laptop?\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m What's the best laptop out there? The Surface Laptop 4 has everything you need.\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m Bye\n",
      "\u001b[1m\n",
      "< / end of conversation. >\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "start_conversation(profile_pic='👦🏻')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "11a61d44-703e-47c7-9803-8296fca021ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m< You are now talking with Mamba. Send 'bye' to exit, 'clear' to reset cache. >\u001b[0m\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m Hi Mamba what's up\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m what's up John, where do you live?\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m I live in New York City. You?\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m I'm from the bay area\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m Nice. Where do you work?\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m I work for this company here in the bay area.\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m What is the company's name?\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m It's a website.\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m Bye\n",
      "\u001b[1m\n",
      "< / end of conversation. >\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "start_conversation(profile_pic='👦🏻')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "304d1330-3808-41d0-b71a-77b62f567404",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[1m< You are now talking with Mamba. Send 'bye' to exit, 'clear' to reset cache. >\u001b[0m\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m Hi Mamba!\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m Hi, John! Have you ever gone skydiving?\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m No I haven't! Have you?\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m Yeah! I went over two months ago and it was really cool! What did you do on your vacation\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m I went to Hawaii! Have you been?\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m Yes, I have!\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m What did you think?\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m It was really beautiful, especially the black sand beaches, they are the best!\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m Awesome! What else did you like?\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m The food! I have never had better seafood in my life!\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m I love seafood too.\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m Have you heard of J.J. Hottie\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m No, who's that...\n",
      "\n",
      "\u001b[48;5;109m\u001b[1m🐍\u001b[0m\u001b[1m Mamba:\u001b[0m J.J. Hottie is a singer, my favorite\n",
      "\n",
      "\u001b[43m\u001b[1m👦🏻\u001b[0m\u001b[1m John:\u001b[0m Bye\n",
      "\u001b[1m\n",
      "< / end of conversation. >\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "start_conversation(profile_pic='👦🏻')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b189ac1d-cbbd-4bdc-88f9-169ee2532afa",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f02f01c-5296-49c3-bda3-c8d124a93e2c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb2522e3-f06f-4b90-ae25-f9dc02da735b",
   "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}