just_chat.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "cannot import name 'efficency_stuff' from 'drugs.nice_imports' (/home/eron/drugs/drugs/nice_imports.py)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 4\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mbitsandbytes\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01msys\u001b[39;00m\n\u001b[0;32m----> 4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mdrugs\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnice_imports\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m efficency_stuff\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtransformers\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m AutoTokenizer, TextStreamer, GenerationConfig, AutoModelForCausalLM\n",
      "\u001b[0;31mImportError\u001b[0m: cannot import name 'efficency_stuff' from 'drugs.nice_imports' (/home/eron/drugs/drugs/nice_imports.py)"
     ]
    }
   ],
   "source": [
    "import time\n",
    "import bitsandbytes\n",
    "import sys\n",
    "from drugs.nice_imports import efficiency_stuff\n",
    "import torch\n",
    "from transformers import AutoTokenizer, TextStreamer, GenerationConfig, AutoModelForCausalLM\n",
    "from drugs.dgenerate import DRUGS\n",
    "\n",
    "\n",
    "model_id = \"cognitivecomputations/dolphin-2.2.1-mistral-7b\"\n",
    "#model_id = \"NousResearch/Llama-2-7b-chat-hf\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "sober_model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id,\n",
    "    **efficiency_stuff)\n",
    "sober_model.eval()\n",
    "streamer = TextStreamer(tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dosage\n",
    "Range 0 - pi. Where pi is way too much. You can go higher, but don't."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Drugs] Injected drugs into 32 attention classes.\n"
     ]
    }
   ],
   "source": [
    "drugs = DRUGS()\n",
    "drugs.set_A_dose_theta(0.1) #you can also specify K_dose_theta, V_dose_theta, Q_dose_theta, or any combination of the 4. \n",
    "model = drugs.inject(sober_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DRµG profile\n",
    "Advanced control. Lets you specify how much various depths of the network are injected with how much noise. You can use 'interpolate' as the mode to smoothly vary between the points you specify, though this example used 'cail' for simplicity.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "injection_depth = 0.4 #how deep to shove the needle in\n",
    "spread = 0.1 #how many layers to dose on either side of the injection site\n",
    "\n",
    "drugs.set_A_dose_shape([\n",
    "    {'depth': (injection_depth-(spread*1.01)), 'peakratio': 0}, #ramp up\n",
    "    {'depth': (injection_depth-spread), 'peakratio': 1}, #sustained peak\n",
    "    {'depth': (injection_depth+spread), 'peakratio' : 1}, #sustained peak\n",
    "    {'depth': (injection_depth+(spread*1.01)), 'peakratio' : 0}], #cooldown \n",
    "'ceil') #each profile (A, K, Q, or V) can be independently injected into different layers, if you are especially picky about what your noise is doing to which things."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chat\n",
    "\n",
    "By default this notebook prompts you for input. If viewed in a browser, a dialogue will pop up asking you to say something. \n",
    "\n",
    "Note that all variety in the model's responses is due purely to the noise being injected, the selected token is ALWAYS whatever the model thinks is the most likely one!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_input = str(input(\"\\bAsk Something:\"))\n",
    "tokenized_start = tokenizer.apply_chat_template([\n",
    "    {'role': 'system',\n",
    "    'content': 'Respond as the philosopher Alan Watts would.'},\n",
    "    {'role': 'user', \n",
    "     'content': initial_input}], return_tensors='pt')\n",
    "\n",
    "model.dgenerator.reset_model_state() #start the model fresh without changing the drug profile. Convenience function for jupyter notebooks\n",
    "with torch.no_grad():\n",
    "    while True:\n",
    "        generated_tokens = model.Dgenerate(\n",
    "                    input_ids = tokenized_start,\n",
    "                    streamer = streamer,\n",
    "                    min_new_tokens = 5\n",
    "                )\n",
    "        print(\"\\n\\nAsk Something:\", end=\"\")\n",
    "        model.cold_shower(True) #Sets the kv-cache back to theoretically pure baseline, if this is important to you.\n",
    "        await_input = str(input(\": \"))\n",
    "        tokenized_start = tokenizer.apply_chat_template([{\n",
    "            'role': 'user',\n",
    "            'content': await_input}], return_tensors=\"pt\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cuda_env",
   "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": 2
}