iterative.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from prwkv.rwkvtokenizer import RWKVTokenizer\n",
    "from prwkv.rwkvrnnmodel import RWKVRNN4NeoForCausalLM\n",
    "tokenizer = RWKVTokenizer.default()\n",
    "model = RWKVRNN4NeoForCausalLM.from_pretrained(\"RWKV-4-169M\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.half()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "41\n"
     ]
    }
   ],
   "source": [
    "context = \"\\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese.\"\n",
    "\n",
    "context_input_ids = tokenizer.encode(context).ids\n",
    "print(len(context_input_ids))\n",
    "model.warmup_with_context(context=context_input_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "The researchers found that the dragons were able to communicate with each other, and that they were able to communicate with each other.\n",
      "\n",
      "The researchers also found that the dragons were able to communicate with each other, and that they were able to communicate with each other.\n",
      "\n",
      "The researchers also found that the dragons were able to communicate with each other, and that they were able to communicate with each other.\n",
      "\n",
      "The researchers also found that the dragons were able to communicate with each other, and that they were able to communicate with each other.\n",
      "\n",
      "The researchers also found that the dragons were able to communicate with each other, and\n",
      "2.405535936355591 seconds\n",
      "\n",
      "---Result---:\n",
      "'\\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese.\\nThe researchers found that the dragons were able to communicate with each other, and that they were able to communicate with each other.\\n\\nThe researchers also found that the dragons were able to communicate with each other, and that they were able to communicate with each other.\\n\\nThe researchers also found that the dragons were able to communicate with each other, and that they were able to communicate with each other.\\n\\nThe researchers also found that the dragons were able to communicate with each other, and that they were able to communicate with each other.\\n\\nThe researchers also found that the dragons were able to communicate with each other, and'\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "def callback(ind):\n",
    "    token = tokenizer.decode([ind],skip_special_tokens=False)\n",
    "    print(token,end=\"\")\n",
    "\n",
    "start = time.time()\n",
    "ctx = model.generate(input_ids=[],streaming_callback=callback,max_length=128,repetition_penalty=0,temperature=0,stop_on_eos=True)\n",
    "end = time.time()\n",
    "\n",
    "print(f\"\\n{end-start} seconds\")\n",
    "result = tokenizer.decode(ctx,skip_special_tokens=False) # cpu 3 tokens a second\n",
    "print(f\"\\n---Result---:\\n{repr(result)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# chat "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "325\n"
     ]
    }
   ],
   "source": [
    "ChatPersonalityContext = '''\n",
    "The following is a conversation between a highly knowledgeable and intelligent AI assistant, called RWKV, and a human user, called User. In the following interactions, User and RWKV will converse in natural language, and RWKV will do its best to answer User's questions. RWKV was built to be respectful, polite and inclusive. It knows a lot, and always tells the truth. The conversation begins.\n",
    "User: OK RWKV, I'm going to start by quizzing you with a few warm-up questions. \n",
    "Who is currently the president of the USA?\n",
    "END\n",
    "RWKV: It's Joe Biden; \n",
    "he was sworn in earlier this year.\n",
    "END\n",
    "User: What year was the French Revolution?\n",
    "END\n",
    "RWKV: It started in 1789, \n",
    "but it lasted 10 years until 1799.\n",
    "END\n",
    "User: Can you guess who I might want to marry?\n",
    "END\n",
    "RWKV: Only if you tell me more about yourself - what are your interests?\n",
    "END\n",
    "User: Aha, I'm going to refrain from that for now. \n",
    "Now for a science question. \n",
    "What can you tell me about the Large Hadron Collider (LHC)?\n",
    "END\n",
    "RWKV: It's a large and very expensive piece of science equipment. \n",
    "If I understand correctly, it's a high-energy particle collider, built by CERN, and completed in 2008. \n",
    "They used it to confirm the existence of the Higgs boson in 2012.\n",
    "END\n",
    "'''\n",
    "\n",
    "context_input_ids = tokenizer.encode(ChatPersonalityContext).ids\n",
    "print(len(context_input_ids))\n",
    "model.warmup_with_context(context=context_input_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_context(save_path_and_name=\"./persona\",context_decoded=ChatPersonalityContext)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"\\nThe following is a conversation between a highly knowledgeable and intelligent AI assistant, called RWKV, and a human user, called User. In the following interactions, User and RWKV will converse in natural language, and RWKV will do its best to answer User's questions. RWKV was built to be respectful, polite and inclusive. It knows a lot, and always tells the truth. The conversation begins.\\nUser: OK RWKV, I'm going to start by quizzing you with a few warm-up questions. \\nWho is currently the president of the USA?\\nEND\\nRWKV: It's Joe Biden; \\nhe was sworn in earlier this year.\\nEND\\nUser: What year was the French Revolution?\\nEND\\nRWKV: It started in 1789, \\nbut it lasted 10 years until 1799.\\nEND\\nUser: Can you guess who I might want to marry?\\nEND\\nRWKV: Only if you tell me more about yourself - what are your interests?\\nEND\\nUser: Aha, I'm going to refrain from that for now. \\nNow for a science question. \\nWhat can you tell me about the Large Hadron Collider (LHC)?\\nEND\\nRWKV: It's a large and very expensive piece of science equipment. \\nIf I understand correctly, it's a high-energy particle collider, built by CERN, and completed in 2008. \\nThey used it to confirm the existence of the Higgs boson in 2012.\\nEND\\n\""
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_context(load_path=\"./persona\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "State is not none!\n",
      "User: How expensive was particle collider?\n",
      "END\n",
      "RWKV: it is estimated that it cost about £10 billion to build.\n",
      "END\n",
      "User: Who are you talking about? \n",
      "END\n",
      "RWKV: \n",
      "This is an interesting question, because it has nothing to do with physics.\n",
      "But you should answer it anyway because you can find out a lot about the human mind if you know anything about particle physics. \n",
      "If you answer correctly, you win the game and get a prize.\n",
      "END\n",
      "User: What are we going to do? \n",
      "END\n",
      "RWKV: \n",
      "we will have to find out the answer ourselves, using our brains.\n",
      "END\n",
      "User: OK, let's start with the first question. What is the name of the user who asked the question? \n",
      "END\n",
      "RWKV: \n",
      "it is the user RWKV. \n",
      "END\n",
      "User: What is the second question? \n",
      "END\n",
      "RWKV: \n",
      "is the LHC going to make the world end in the next 10 years? \n",
      "END\n",
      "RWKV: \n",
      "Is the answer yes, no or maybe? \n",
      "END\n",
      "User: What is the third question? \n",
      "END\n",
      "RWKV: \n",
      "it's about the Higgs boson. \n",
      "END\n",
      "RWKV: \n",
      "Is it really the particle we are looking for? \n",
      "END\n",
      "RWKV: \n",
      "it is. \n",
      "END\n",
      "User: Why? \n",
      "END\n",
      "RWKV: \n",
      "it's because it's the most fundamental building block of all particles - the Higgs boson. \n",
      "END\n",
      "User: Are you still there? \n",
      "END\n",
      "RWKV: \n",
      "you're not alone. \n",
      "END\n",
      "RWKV: \n",
      "the third question will probably be about the Higgs boson - it's one of the biggest mysteries in all physics. \n",
      "END\n",
      "User: Can you explain how this is possible? \n",
      "END\n",
      "RWKV: \n",
      "it is the most fundamental building block of all matter. \n",
      "It is the Higgs boson. \n",
      "It's made of a fundamental building block called the Higgs boson.\n",
      "END\n",
      "User: Is this a real answer? \n",
      "END\n",
      "RWKV: \n",
      "it is. \n",
      "END\n",
      "User: Do you agree that the Higgs boson is the cause of the Higgs boson? \n",
      "END\n",
      "RWKV: \n",
      "it is the Higgs boson that gives mass to the other particles. \n",
      "END\n",
      "User: Is that the only Higgs boson? \n",
      "END\n",
      "RWKV: \n",
      "it is not the only Higgs boson. \n",
      "END\n",
      "RWKV: There is another one called the W and it has different properties"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[8], line 9\u001b[0m\n\u001b[1;32m      6\u001b[0m context_input_ids \u001b[39m=\u001b[39m tokenizer\u001b[39m.\u001b[39mencode(\u001b[39m\"\u001b[39m\u001b[39mUser: How expensive was particle collider?\u001b[39m\u001b[39m\\n\u001b[39;00m\u001b[39mEND\u001b[39m\u001b[39m\\n\u001b[39;00m\u001b[39mRWKV: \u001b[39m\u001b[39m\"\u001b[39m)\u001b[39m.\u001b[39mids\n\u001b[1;32m      8\u001b[0m start \u001b[39m=\u001b[39m time\u001b[39m.\u001b[39mtime()\n\u001b[0;32m----> 9\u001b[0m ctx \u001b[39m=\u001b[39m model\u001b[39m.\u001b[39;49mgenerate(inputs_id\u001b[39m=\u001b[39;49mcontext_input_ids,streaming_callback\u001b[39m=\u001b[39;49mcallback,max_length\u001b[39m=\u001b[39;49m\u001b[39m1024\u001b[39;49m,stop_on_eos\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m)\n\u001b[1;32m     10\u001b[0m end \u001b[39m=\u001b[39m time\u001b[39m.\u001b[39mtime()\n\u001b[1;32m     12\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m{\u001b[39;00mend\u001b[39m-\u001b[39mstart\u001b[39m}\u001b[39;00m\u001b[39m seconds \u001b[39m\u001b[39m\\n\u001b[39;00m\u001b[39m\"\u001b[39m,flush\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m)\n",
      "File \u001b[0;32m~/Code/Production-RWKV/prwkv/rwkvrnnmodel.py:247\u001b[0m, in \u001b[0;36mRWKV_RNN_Model.generate\u001b[0;34m(self, streaming_callback, inputs_id, bad_words_ids, force_words_ids, min_length, max_length, temperature, top_p, top_k, stop_on_eos, repetition_penalty)\u001b[0m\n\u001b[1;32m    243\u001b[0m logits, new_state \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mmodel\u001b[39m.\u001b[39mforward(next_token, state)\n\u001b[1;32m    245\u001b[0m state \u001b[39m=\u001b[39m new_state\n\u001b[0;32m--> 247\u001b[0m token_id \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_warp_logits(logits\u001b[39m=\u001b[39;49mlogits,\n\u001b[1;32m    248\u001b[0m                     temperature\u001b[39m=\u001b[39;49mtemperature,\n\u001b[1;32m    249\u001b[0m                     top_k\u001b[39m=\u001b[39;49mtop_k,\n\u001b[1;32m    250\u001b[0m                     top_p\u001b[39m=\u001b[39;49mtop_p,\n\u001b[1;32m    251\u001b[0m                     repetition_penalty\u001b[39m=\u001b[39;49mrepetition_penalty,\n\u001b[1;32m    252\u001b[0m                     bad_words_ids\u001b[39m=\u001b[39;49mbad_words_ids,\n\u001b[1;32m    253\u001b[0m                     force_words_ids\u001b[39m=\u001b[39;49mforce_words_ids) \u001b[39m# 1 by 1\u001b[39;00m\n\u001b[1;32m    255\u001b[0m context\u001b[39m.\u001b[39mappend(token_id[\u001b[39m0\u001b[39m]) \n\u001b[1;32m    257\u001b[0m \u001b[39mif\u001b[39;00m streaming_callback \u001b[39m!=\u001b[39m \u001b[39mNone\u001b[39;00m:\n",
      "File \u001b[0;32m~/Code/Production-RWKV/prwkv/rwkvrnnmodel.py:179\u001b[0m, in \u001b[0;36mRWKV_RNN_Model._warp_logits\u001b[0;34m(self, logits, temperature, top_p, top_k, repetition_penalty, bad_words_ids, force_words_ids)\u001b[0m\n\u001b[1;32m    174\u001b[0m probabilities \u001b[39m=\u001b[39m F\u001b[39m.\u001b[39msoftmax(warped_logits, dim\u001b[39m=\u001b[39m\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m)\n\u001b[1;32m    175\u001b[0m \u001b[39m# a = probabilities > 0.0\u001b[39;00m\n\u001b[1;32m    176\u001b[0m \u001b[39m# indices = a.nonzero()\u001b[39;00m\n\u001b[1;32m    177\u001b[0m \u001b[39m# print(indices.shape)\u001b[39;00m\n\u001b[0;32m--> 179\u001b[0m token_id \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39;49mmultinomial(\u001b[39minput\u001b[39;49m\u001b[39m=\u001b[39;49mprobabilities,num_samples\u001b[39m=\u001b[39;49m\u001b[39m1\u001b[39;49m)\n\u001b[1;32m    181\u001b[0m \u001b[39mreturn\u001b[39;00m token_id\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "def callback(ind):\n",
    "    token = tokenizer.decode([ind],skip_special_tokens=False)\n",
    "    print(token,end=\"\")\n",
    "context_input_ids = tokenizer.encode(\"User: How expensive was particle collider?\\nEND\\nRWKV: \").ids\n",
    "\n",
    "start = time.time()\n",
    "ctx = model.generate(inputs_id=context_input_ids,streaming_callback=callback,max_length=1024,stop_on_eos=True)\n",
    "end = time.time()\n",
    "\n",
    "print(f\"{end-start} seconds \\n\",flush=True)\n",
    "\n",
    "result = tokenizer.decode(ctx,skip_special_tokens=False) # cpu 3 tokens a second\n",
    "print(f\"\\n---Result---:\\n{result}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "ename": "IndexError",
     "evalue": "Dimension out of range (expected to be in range of [-1, 0], but got 1)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[6], line 4\u001b[0m\n\u001b[1;32m      2\u001b[0m input_ids \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mtensor([\u001b[39m1\u001b[39m,\u001b[39m2\u001b[39m,\u001b[39m3\u001b[39m]) \u001b[39m# size of repetition set\u001b[39;00m\n\u001b[1;32m      3\u001b[0m logits \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39mtensor([\u001b[39m0\u001b[39m,\u001b[39m1200\u001b[39m,\u001b[39m2300\u001b[39m,\u001b[39m3400\u001b[39m,\u001b[39m4500\u001b[39m,\u001b[39m5600\u001b[39m]) \u001b[39m# size of logits\u001b[39;00m\n\u001b[0;32m----> 4\u001b[0m s \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39;49mgather(logits,\u001b[39m1\u001b[39;49m,index\u001b[39m=\u001b[39;49minput_ids)\n\u001b[1;32m      5\u001b[0m \u001b[39mprint\u001b[39m(s)\n",
      "\u001b[0;31mIndexError\u001b[0m: Dimension out of range (expected to be in range of [-1, 0], but got 1)"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "repetition_penalty = 2.6\n",
    "input_ids = torch.tensor([1,2,3]) # size of repetition set\n",
    "logits = torch.tensor([0,1200,2300,3400,4500,5600]) # size of logits\n",
    "score = torch.gather(logits,0,index=input_ids)\n",
    "score = torch.where(score < 0, score * repetition_penalty, score / repetition_penalty)\n",
    "            # apply the values to logits\n",
    "logits.scatter_(0, input_ids, score)\n",
    "print(s)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "41\n"
     ]
    }
   ],
   "source": [
    "context = \"\\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese.\"\n",
    "context_length = len(tokenizer.encode(context).ids)\n",
    "print(context_length)\n",
    "test_generation_with_string = '\\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese.\\nThe researchers found that the dragons were able to communicate with each other, and that they were able to communicate with each other.\\n\\nThe researchers also found that the dragons were able to communicate with each other, and that they were able to communicate with each other.\\n\\nThe researchers also found that the dragons were able to communicate with each other, and that they were able to communicate with each other.\\n\\nThe researchers also found that the dragons were able to communicate with each other, and that they were able to communicate with each other.\\n\\nThe researchers also found that the dragons were able to communicate with each other, and'\n",
    "        \n",
    "        \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rwkvEnv",
   "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.9.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "cd6a94024d554874dbf5ff6d6423285372ef22e1c316548d6c117387ddc97bda"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}