rope_test.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Following code is simplified 4-bit model loader to demonstrate that any model using Rotary Positional Embeddings (i.e, LLaMa) and finetuned with X cutoff len has X context limit\n",
    "\n",
    "Context limit is entirely dependent on the positional embedding scheme. With RoPE, the model is generalizable to any length, however the extrapolation is not very good due to the destabilizing issues in default RoPE implementation. This is why you experience gibberish text when trying to use a Alpaca model (512 cutoff len) past 2048. The reason Alpaca still works well beyond 512 tokens is because the learned positional embeddings weights beyond 512 are still retained (base LLaMa is trained on 2048 after all) Bluemoon (4096 cutoff) also works well past 2048 for the same reason, although due to the training it doesn't work quite as well (the scale of finetuning is nowhere near pre-training)\n",
    "\n",
    "The extrapolation issue can be fixed by using the implementation described in https://arxiv.org/abs/2212.10554v1 (\"A Length-Extrapolatable Transformer\") and unlike ALiBi is compatible with RoPE (as it is a generalization of the latter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from util.amp_wrapper import AMPWrapper\n",
    "from util.autograd_4bit import load_llama_model_4bit_low_ram\n",
    "from util.peft_tuners_lora_monkey_patch import replace_peft_model_with_gptq_lora_model\n",
    "replace_peft_model_with_gptq_lora_model()\n",
    "import torch\n",
    "from tqdm import tqdm\n",
    "from peft import PeftModel"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_safetensors_path = \"/.../.../...safetensors\"\n",
    "model_dir = \"/.../\"\n",
    "lora_dir = \"/.../\"\n",
    "groupsize = -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model, tokenizer = load_llama_model_4bit_low_ram(\n",
    "    model_dir,\n",
    "    model_safetensors_path,\n",
    "    half=True,\n",
    "    is_v1_model=False,\n",
    "    groupsize=groupsize,\n",
    ")\n",
    "wrapper = AMPWrapper(model)\n",
    "wrapper.apply_generate()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optionally apply LoRA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "PeftModel.from_pretrained(model, lora_dir, dtype=model.dtype)\n",
    "\n",
    "model.half()\n",
    "model = model.cuda()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Write a prompt. To test, you should put a password in the beginning or end of the string to see that is able to fetch it either way\n",
    "\n",
    "The default prompt here puts the password in the beginning and asks for it at the end while priming with a single letter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "password = \"BLUEMOON\"\n",
    "sample = \"\"\"The password was {password}. {password} is the password. Soon the time will come when we must bid adieu to our beloved homeland. When that day comes, you and I shall have to go our separate ways. But before that happens, I want to make sure that you are safe and well. To ensure this, I would like you to travel to Fanelia, where you will be enrolled in a prestigious school. I am certain that you will do very well there, and that the Headmaster and staff will take good care of you. And once you have finished your studies, I am sure that many doors will be opened for you.' She smiled at him again, but this time her eyes were glistening with unshed tears. 'Now, my dear child, I must away. But before I leave, I want you to give me your word that you will do as I have asked.'\n",
    "Justin stared up at Alenia, unable to say a word. He knew that if he refused her, she would simply hire someone else to take him to Fanelia, but he had never dreamed that she would just leave him behind! 'Alenia...' he finally choked out, his voice wavering. 'I... I can't go to Fanelia.'\n",
    "'Oh?' Alenia looked a little taken aback. 'And why not, may I ask?'\n",
    "Justin looked down at the ground and fidgeted with his shoes for a moment before he looked up at her again. 'Because I... I want to stay here in Domino and help you fight the war.' Now that the words were out, he felt an overwhelming sense of relief, as though a great weight had been lifted from his shoulders.\n",
    "Alenia's face paled somewhat and she looked very troubled. 'Justin,' she said gently, 'I know you are a good boy, but this is something that you cannot understand. You will be of no use to me or Fanelia if you throw your life away. I must insist that you go to Fanelia.'\n",
    "'I... I can't let you do that, Alenia!' he protested. 'You've been like a mother to me, and I couldn't stand it if anything happened to you!' His eyes were pleading now, and he desperately wanted her to change her mind.\n",
    "Alenia sighed and gently stroked his cheek. 'Then I shall make you a promise in return,' she said. 'I will come and visit you as often as I can, and if the war should end during your time at school...' She smiled at him, but it was a sad smile. 'Then you can return to Domino and I will send you to the best school here in Hibiki.'\n",
    "Justin opened his mouth to say something, but Alenia placed a finger to his lips. 'No, let me finish,' she said calmly. 'I shall not allow you to throw your life away, Justin. So you see, it is better this way. Now won't you do as I have asked? I must go and prepare for my journey.'\n",
    "And with that, Alenia stood up and gathered her cloak about her. 'It grows cold here,' she said. 'So don't forget your coat on the way to the carriage.' With those words, she swept out of the room without looking back.\n",
    "Justin swallowed the lump in his throat and got to his feet. Somehow he could not bring himself to protest any further. If he did, then Alenia would hire someone else to take him to Fanelia, and he would be forced to go anyway. And although that thought was preferable to being left behind in Domino, it still wasn't what he wanted. What he wanted was for Alenia to change her mind... but she never did.\n",
    "Once or twice in the days that followed, Justin considered sneaking off to try and find Alenia, but each time he was about to do so, he remembered her words: 'Now don't you dare go after me, you stubborn boy.' She had known him long enough to know that if he really wanted to do something, nothing would stop him. The only reason he hadn't tried to find her was because... he didn't want to go to Fanelia.\n",
    "The day he was to leave for Fanelia arrived all too soon for Justin's liking, and when the time came for him to bid Alenia farewell, he felt anguish like he never had before. They had stopped at a small rest area where travellers could stop to eat or sleep before continuing on their journey, and Alenia was waiting for him there with her carriage. She was wearing a beautiful red dress, and she looked very beautiful indeed. But all Justin could see was the misery in her eyes, and his heart ached.\n",
    "Alenia came over to him as he got down from his horse and pulled him into her arms. 'Justin,' she murmured, 'I will miss you so.' Then she sighed and gently stroked his cheek. 'Be good, my dear child... and study well.' Her voice broke a little then, and she pulled away from him to get back up onto the carriage. There was one last thing she had to do before they parted... She had to make sure that he would go to Fanelia. So with aching heart, Alenia took out a small leather case and drew out a hypodermic needle. Without even looking at it, she knew that it contained a truth serum of some kind. All she had to do was prick Justin's neck with it, and he would be compelled to do her bidding. But just as she was about to do so, she hesitated. A tiny voice in the back of her head had begun to whisper that this was wrong... that she was taking away his free will. And since when did she start caring about such things? The voice was quickly silenced, however, and with a determined set to her lips, Alenia steeled herself to do what she had to do. She would not be swayed from her decision. So with a heavy heart, she stabbed Justin's neck and then jumped down from the carriage to join her travelling companions in the back. They started up the carriage immediately, and Alenia sat with her head bowed, weeping silently for the child she was losing. And all because he didn't want to go to Fanelia...\n",
    "The journey to Fanelia was a long one, and boring too. After the first few hours, Justin had been longing to get out of the carriage and stretch his legs, but Alenia's instructions had been very specific: no stops unless absolutely necessary. The strict sense of discipline in which he had been raised rebelled at the thought of not being able to do as he pleased for so long, but he knew that if he did, then Alenia would hire someone else to take him to Fanelia, and he certainly didn't want that. So he gritted his teeth and tried to ignore the urge to move about. Several times he was nearly struck by the driver's whip as it cracked overhead, but by and large, he was a good traveller and did not stir from his seat.\n",
    "Justin's only consolation was that they were nearing Fanelia. The feeling of excitement at the thought was short lived, however, for as they entered the city, he felt a strange sensation... Alenia had used truth serum on him! He knew it! But why? And what had she said to him? Justin racked his brain, trying desperately to remember... But all that happened was an overwhelming sense of confusion. It was like his mind had been... scrubbed clean. All he could really remember was Alenia's face and her voice. Everything else was just a jumbled mess in his head, and no matter how hard he tried, nothing would come to the surface. Even now, as he attempted to write this journal entry, it felt as though there was a fog obscuring his memory...\n",
    "Justin sighed and set down his quill. It didn't seem fair that Alenia could make him forget whatever she had told him. But even if he did remember, what could he do about it? He was in Fanelia now, and there wasn't a thing he could do to change that. Perhaps it was for the best... After all, what did he really know about Fanelia aside from the fact that they were involved in some kind of war? Not much, as it turned out. But Justin knew what he wanted to do. He wanted to help Alenia and fight in the war. And if that meant having a truth serum administered to him every now and then... Well, he could just ask Alenia about it when he next saw her. Which would be when, exactly? That was another thing he didn't know... But surely he would be able to see her again sometime... Someone had mentioned something about a school, but Justin hadn't been paying much attention. All he wanted to do at that moment was find the inn where he was to stay the night and get some sleep. Wearily, he rose to his feet and made from the tavern in which he had been writing to the inn where he would be spending the night. He couldn't remember the name of it, but it didn't really matter. All that mattered was that he got some rest. He had a feeling he would need it for the journey ahead... But what journey? There was still so much Justin did not know...\n",
    "Later that evening, when Justin had gone to bed in his room at the inn, he found himself remembering Alenia's parting words. 'Study well,' she had said. What did she mean by that? He had never been to school before, so why would he need to study? Justin frowned and tossed and turned for a long time in his bed, but sleep would not come to him. Giving up, he got out of bed and went downstairs to see if the innkeeper would give him something to drink. Perhaps it would help him to relax and make him drowsy so that he could finally get some rest. As he walked down the stairs, however, he heard a most curious sound... The sounds of two people arguing, it seemed. Interest piqued, Justin headed over to the pair and listened in on their conversation. From what he had gathered, they were a father and son who had business in Fanelia. The father's name was Richard, and he was an archaeologist of some kind. He and his son, Dmitri, were on a special mission for the Domino Empire. They were on their way to see the Director of the Fanelian School of Arts and Letters, who would be taking Justin as well. But Justin's part in this affair was still vague. All he knew was that he was to be enrolled in the school and that they would be taking you away from here.' That's when it comes, Justin thought with a small sigh. 'But I am not ready to leave Alenia.' A shadow of confusion crossed his face, and a little bit of fear as well. If he is not ready to leave Alenia, then who is? Will no one will be there to see me off. The realisation that he was in Fanelia for good and never look back, for no matter what happens, I must do my best to make Alenia proud of me and fight in the war to the best of my ability. A small tear slipped down Justin's cheek and splashed upon the floor. It wasn't the first time he had cried, and surely it wouldn't be the last. Once again, he will be alone in a new country, with no place to call home. But if Fanelia was to become his new home, then he would have to make new friends and live a life worthy of Alenia's fondness and faith in him. That thought brought a faint smile to his face and a surge of confidence that somehow, no matter how far away she was, Alenia would watch over him and protect him as she had before. But why would she bother to do so? He was old enough and wise enough to take care of himself. And with that, Justin was no longer a child, but a young man, finally grown up completely and ready to embrace his new life and the many friends he will make there. And although part of him knew that he would be alone in this journey, it wouldn't be the last time he would be alone. And until then... He would continue onwards as an adventurer and try my best to be worthy of Alenia's pride and faith in me. I must study hard and achieve a scholarship to Alenia and make her proud of me, but also protecting her as she travels to and fro from Hibiki and eventually, to Fanelia. Surely it wouldn't be long before he would not see her again for a while and probably, forever. And even if it does, I will always remember her and the days we spent together. Often, when I do, I will think of you and smile, Alenia, because no matter what happens, I will always keep you close to me in my heart and in my memories. So goodbye, my dearest girl, and God speed you, Alenia! And with that, Justin's journey to Fanelia commenced... First, she remembered the password: {password[0]}\"\"\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the token size of context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2943"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(tokenizer.encode(sample, return_tensors='pt')[0])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate\n",
    "\n",
    "Any model trained with > 2048 cutoff should be able to print the correct password from above. \n",
    "Due to the stability issue in RoPE, it may need priming, but it can definitely get the password\n",
    "\n",
    "TL;DR you should train with higher cutoff if you have the data for it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "token_len = len(tokenizer.encode(sample, return_tensors='pt')[0])\n",
    "print(\"Token context len:\", token_len)\n",
    "tokenizer.decode(model.generate(tokenizer.encode(sample, return_tensors='pt'), max_new_tokens=50)[0][token_len-1:])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "textgen",
   "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.10.9"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}