testing_minichat.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import peft \n",
    "import time\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "from utils.dataset import template\n",
    "\n",
    "#from bokeh.plotting import figure, show\n",
    "#from bokeh.io import output_notebook\n",
    "\n",
    "#output_notebook()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = 'GeneZC/MiniChat-3B'\n",
    "device = 'cuda'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using the default legacy behaviour of the <class 'transformers.models.llama.tokenization_llama.LlamaTokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thouroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565\n"
     ]
    }
   ],
   "source": [
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.bfloat16\n",
    ")\n",
    "\n",
    "\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(model_id, use_cache=True,\n",
    "                                             device_map=\"auto\",\n",
    "                                             quantization_config=bnb_config,\n",
    "                                             )\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=False)\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.add_special_tokens({'pad_token': '[PAD]'})\n",
    "    model.resize_token_embeddings(len(tokenizer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "lora_path = 'lora/disco-limbic-dialogue-512/'\n",
    "\n",
    "model = model.eval()\n",
    "lora_model = peft.PeftModel.from_pretrained(model, lora_path, adapter_name='loraTrained', is_trainable=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " You are the parts of the human brain that conduct a dialogue, you can enter into verbal altercations with the interlocutor. You need to response emotionally.[|Assistant|] [Electrochemistry]: Whoa! In your hand: *pyrholidon* -- the double rainbow of synthetic hallucinogens. Rare and gritty, a product of the age of atomic power.</s> [|User|] Look at the little puck of liquid.</s> [|Assistant|] [Electrochemistry]: What a funny little cap! Don't let the *scary* medical warnings throw you off. It's an inadequate antidote to radiation poisoning, but a *potent* antidote to *boredom*.</s> [|User|] Hmm... open the cap.</s>  [|Assistant|]  [Electrochemistry]: The smell is acrid, sharp, almost chemical on its own.\n",
      "[Sight]: Go ahead and sniff it, see how hard it comes back? How long does the after-effect last?</s>\n"
     ]
    }
   ],
   "source": [
    "dialog = [\n",
    "    \"[Electrochemistry]: Whoa! In your hand: *pyrholidon* -- the double rainbow of synthetic hallucinogens. Rare and gritty, a product of the age of atomic power.\",\n",
    "    \"Look at the little puck of liquid.\",\n",
    "    \"[Electrochemistry]: What a funny little cap! Don't let the *scary* medical warnings throw you off. It's an inadequate antidote to radiation poisoning, but a *potent* antidote to *boredom*.\",\n",
    "    \"Hmm... open the cap.\",\n",
    "]\n",
    "\n",
    "query = template(dialog) + ' [|Assistant|] '\n",
    "model_inputs = tokenizer(query, return_tensors=\"pt\", add_special_tokens=False).input_ids.to(device)\n",
    "generated_ids = model.generate(input_ids=model_inputs, max_new_tokens=64,\n",
    "                               do_sample=True,\n",
    "                               temperature=0.7,\n",
    "                               repetition_penalty=1.15)\n",
    "\n",
    "output = tokenizer.decode(generated_ids[0], skip_special_tokens=False)\n",
    "print(output)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Interfacing]: The pink and white ball slips out of your hand as it hits the ice -- what a *punk* moment...\n",
      "\n",
      "[Visual Calculus]: And just look at that! A communist symbol in all its glory. It's like an army flag or something.\n",
      "[Visual Calculus]: What is this crap? Are we really so lousy? Why isn't it anything other than a piece of shit?\n",
      "\n"
     ]
    }
   ],
   "source": [
    "dialog = [\n",
    "]\n",
    "\n",
    "while True:\n",
    "    \n",
    "    inp = input('Input:')\n",
    "    if inp == 'q':\n",
    "        break\n",
    "    if inp:\n",
    "        dialog.append(inp)\n",
    "        \n",
    "        print(f'[You]: {inp}', end='\\n\\n')\n",
    "    query = template(dialog) + ' [|Assistant|] '\n",
    "    \n",
    "    model_inputs = tokenizer(query, return_tensors=\"pt\", add_special_tokens=False).input_ids.to(device)\n",
    "    input_len = len(model_inputs[0])\n",
    "    generated_ids = model.generate(input_ids=model_inputs, max_new_tokens=128,\n",
    "                                do_sample=True,\n",
    "                                #pad_token_id=tokenizer.eos_token_id,\n",
    "                                temperature=0.7,\n",
    "                                repetition_penalty=1.15)\n",
    "    #outputs = model(input_ids=input_ids),# max_length=cut_len, min_length=8, top_p=0.9, do_sample=True)\n",
    "    output = tokenizer.decode(generated_ids[0][input_len:], skip_special_tokens=True)\n",
    "    dialog.append(output)\n",
    "    print(output, end='\\n\\n')\n",
    "    time.sleep(0.5)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "disco",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}