textual_avatar.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b89eafb3",
   "metadata": {},
   "source": [
    "# Textual Avatar - Train LLM on your own!\n",
    "\n",
    "This notebook intend to export your dialogs into encrpyted instruction based format and fine-tune LLama-2 on it.\n",
    "It is also provided gradio demo and Telegram bot with inferencing model.\n",
    "For fine tune LLama-2 you need access to capable hardware, I use Nvidia RTX-3090 for one night.\n",
    "\n",
    "## Content\n",
    "\n",
    "In this project, we'll follow these main steps:\n",
    "\n",
    "\n",
    "### 1. [Export Meta Messenger data](#Export-Meta-Messenger-data) or \n",
    "### 2. [Export Telegram data](#Exporting-from-Telegram) or\n",
    "### 3. (Optionally) [Export other data in instruction-based format](#Add-more-data)\n",
    "### 4. [Preprocess and encrypt dataset file](#Preprocess-and-encrypt-dataset-file)\n",
    "### 5. [Download LLM and setup machine which support fine-tune LLM](#Setup-transformers,-PEFT-and-LLM)\n",
    "### 6. [Load dataset and fine-tune LLM for several hours](#Training-script,-simplified-from-HF-tutorial)\n",
    "### 7. [Run chat with yourself](#Compare-with-previous-replies)\n",
    "### 8. [Create chat bot](#Telegram-chat-bot)\n",
    "\n",
    "\n",
    "## Privacy and Security\n",
    "\n",
    "I provide encryption of data file, all private files could be deleted after preprocessing.\n",
    "\n",
    "However, exercise caution and only run this notebook on trusted systems.\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0ef569e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import glob\n",
    "import json\n",
    "import os\n",
    "import pickle\n",
    "import csv\n",
    "\n",
    "import pathlib\n",
    "from getpass import getpass\n",
    "\n",
    "'''\n",
    "Only export these messages\n",
    "'''\n",
    "\n",
    "\n",
    "def filter_message(msg: str) -> bool:\n",
    "    return len(msg) > 1 and not msg.startswith(\"http\")\n",
    "\n",
    "\n",
    "TEXTUAL_AVATAR = \"<TEXTUAL_AVATAR>\"\n",
    "\n",
    "'''\n",
    "Instruction for LLM, using previous conversation\n",
    "'''\n",
    "\n",
    "\n",
    "def convert_sample_to_instruction(sample):\n",
    "    previous_conversation = \"\"\n",
    "    for message in sample['last_messages']:\n",
    "        previous_conversation = previous_conversation + \\\n",
    "            message[\"author\"]+\":\"+message[\"text\"]+\"\\n\"\n",
    "    previous_conversation = previous_conversation[-4096:]\n",
    "\n",
    "    instruction = f\"\"\"You are {TEXTUAL_AVATAR}, a sophisticated AI designed to engage in text conversations.\n",
    "    Your goal is to provide relevant responses based on the given context.\n",
    "    Imagine you have been having a conversation with {sample['counterpart']}\n",
    "    Your task is to mimic a text reply to last message as {TEXTUAL_AVATAR}\n",
    "    \"\"\"\n",
    "    input = f\"\"\"previous conversation was:\\n{previous_conversation} last message is {sample['input']}\"\"\"\n",
    "    response = sample['text']\n",
    "    return {\"instruction\": instruction, \"input\": input, \"response\": response}\n",
    "\n",
    "\n",
    "MAX_PREVIOUS_MESSAGES = 12\n",
    "KEEP_CONVERSATION_SEC = 60*60*24\n",
    "\n",
    "with open('textual_avatar.csv', 'w', newline='') as csvfile:\n",
    "    writer = csv.DictWriter(csvfile, fieldnames=[\n",
    "                            \"instruction\", \"input\", \"response\"])\n",
    "    writer.writeheader()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bdd6ee0",
   "metadata": {},
   "source": [
    "# Export Meta Messenger data\n",
    "\n",
    "To download messages from Facebook:\n",
    "1. Click on Setting->Your Facebook Information -> Download your information https://www.facebook.com/dyi/\n",
    "2. Choose your messages, for period **All Time**, format JSON\n",
    "3. Wait (1 hour) until data collected, download it and unzip\n",
    "4. Uncomment line 5 and 6 in next cell with  ```YOUR_PATH``` to unzipped folder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52064dd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "exported_fb_file_paths = glob.glob(\n",
    "    f\"{os.environ['HOME']}/Downloads/messages/inbox/*/message*.json\")\n",
    "\n",
    "# or enter here your location for FB archive with jsons\n",
    "# YOUR_PATH=\n",
    "# exported_fb_file_paths=glob.glob(YOUR_PATH+\"/messages/inbox/*/message*.json\")\n",
    "\n",
    "\n",
    "encoding = 'latin1'\n",
    "participants = {}\n",
    "for json_file in exported_fb_file_paths:\n",
    "    with open(json_file, 'r', encoding='utf-8') as file:\n",
    "        data = json.load(file)\n",
    "        if len(data['messages']) > 2 and len(data['participants']) == 2:\n",
    "            for participant in data[\"participants\"]:\n",
    "                name = participant[\"name\"].encode(encoding).decode()\n",
    "                if name not in participants.keys():\n",
    "                    participants[name] = set()\n",
    "                participants[name].add(json_file)\n",
    "\n",
    "max_length = 0\n",
    "you = \"\"\n",
    "for key, value in participants.items():\n",
    "    if len(value) > max_length:\n",
    "        you = key\n",
    "        max_length = len(value)\n",
    "\n",
    "\n",
    "chats = {}\n",
    "your_messages_idx = {}\n",
    "for json_file in exported_fb_file_paths:\n",
    "    with open(json_file, 'r', encoding='utf-8') as file:\n",
    "        data = json.load(file)\n",
    "        chat = []\n",
    "        yours = []\n",
    "        for participant in data[\"participants\"]:\n",
    "            name = participant[\"name\"].encode(encoding).decode()\n",
    "            if len(name) < 1:\n",
    "                name = \"Facebook user\"\n",
    "            if name != you:\n",
    "                participant_name = name\n",
    "        if len(data['messages']) > 2 and len(data['participants']) == 2:\n",
    "            for i, message in enumerate(data['messages']):\n",
    "                if 'content' in message.keys():\n",
    "                    author = message['sender_name'].encode(encoding).decode()\n",
    "                    if author == you:\n",
    "                        yours.append(i)\n",
    "                        author = \"TEXTUAL_AVATAR\"\n",
    "                    text = message['content'].encode(encoding).decode()\n",
    "                    if filter_message(text):\n",
    "                        new_message = {\"text\": text,\n",
    "                                       \"tstamp\": message['timestamp_ms']//1000,\n",
    "                                       \"author\": author}\n",
    "                        chat.append(new_message)\n",
    "            chats[participant_name] = chat\n",
    "            your_messages_idx[participant_name] = yours\n",
    "\n",
    "\n",
    "print(f\"exporting {len(chats)} chats\")\n",
    "\n",
    "with open('text_avatar.csv', 'w', newline='') as csvfile:\n",
    "    writer = csv.DictWriter(csvfile, fieldnames=[\n",
    "                            \"instruction\", \"input\", \"response\"])\n",
    "    for counterpart, chat in chats.items():\n",
    "        for your_reply in your_messages_idx[counterpart]:\n",
    "            last_message_to_use = max(your_reply-MAX_PREVIOUS_MESSAGES, 0)\n",
    "            your_reply_tstamp = chat[your_reply][\"tstamp\"]\n",
    "\n",
    "            for message_idx in reversed(range(last_message_to_use, your_reply)):\n",
    "                if (your_reply_tstamp-chat[message_idx][\"tstamp\"]) > KEEP_CONVERSATION_SEC:\n",
    "                    last_message_to_use = message_idx+1\n",
    "                    break\n",
    "            sample = {\"counterpart\": counterpart,\n",
    "                      \"last_messages\": chat[last_message_to_use:your_reply-1],\n",
    "                      \"input\": chat[your_reply-1][\"text\"],\n",
    "                      \"text\": chat[your_reply][\"text\"]}\n",
    "            writer.writerow(convert_sample_to_instruction(\n",
    "                sample))  # adding a row\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17429389",
   "metadata": {},
   "source": [
    "# Exporting from Telegram\n",
    "\n",
    "To collect messages from Telegram on Mac I did the following:\n",
    "\n",
    "1. Download the official Telegram client on Mac from here https://macos.telegram.org  (Mac App Store version not worked for me)\n",
    "2. Click Setting->Advanced->Export Telegram data on the bottom. Current script work with personal chat and machine friendly JSON\n",
    "3. Confirm this action in chat on other device.\n",
    "4. Wait until download process finish. Edit line 4 below ```exported_tg_file_path``` with your path to result.json\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "a0bae59e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "found 236 chats\n"
     ]
    }
   ],
   "source": [
    "#exported_tg_file_path = glob.glob(\n",
    "#    f\"{os.environ['HOME']}/Downloads/Telegram Desktop/DataExport_*/result.json\")[0]\n",
    "# or enter here your location for\n",
    "exported_tg_file_path =\"result.json\"\n",
    "with open(exported_tg_file_path) as exported_tg_file:\n",
    "    exported = json.load(exported_tg_file)\n",
    "chats = [chat for chat in exported[\"chats\"]\n",
    "         [\"list\"] if chat['type'] == 'personal_chat']\n",
    "print(f\"found {len(chats)} chats\")\n",
    "\n",
    "participants = {}\n",
    "participant_names = {}\n",
    "you = \"\"  # tricky way to find your id assume you appear in most chats\n",
    "\n",
    "for chat in chats:\n",
    "    current_chat = {}\n",
    "    for message in chat['messages']:\n",
    "        if message['type'] == \"message\":\n",
    "            if 'from_id' in message.keys():\n",
    "                if message['from_id'] not in participants.keys():\n",
    "                    participants[message['from_id']] = set()\n",
    "                participants[message['from_id']].add(chat['id'])\n",
    "                participant_names[message['from_id']] = message['from']\n",
    "\n",
    "\n",
    "max_length = 0\n",
    "you = \"\"\n",
    "for key, value in participants.items():\n",
    "    if len(value) > max_length:\n",
    "        you = key\n",
    "        max_length = len(value)\n",
    "\n",
    "exported_chats = {}\n",
    "your_messages_idx = {}\n",
    "\n",
    "for conversation in chats:\n",
    "    idx = 0\n",
    "    chat = []\n",
    "    yours = []\n",
    "    for message in conversation['messages']:  # pick counterpart user name\n",
    "        if message['type'] == \"message\":\n",
    "            if message[\"from_id\"] != you:\n",
    "                participant_name = message[\"from\"]\n",
    "                break\n",
    "\n",
    "    for message in conversation['messages']:\n",
    "        if message['type'] == \"message\" and \"from_id\" in message.keys() and filter_message(str(message[\"text\"])):\n",
    "            text = \"\"\n",
    "            if isinstance(message[\"text\"], list):\n",
    "                for text_item in message[\"text\"]:\n",
    "                    if isinstance(text_item, str):\n",
    "                        text = text+\"\\n\"+text_item\n",
    "            else:\n",
    "                text = message[\"text\"]\n",
    "            if message[\"from_id\"] == you:\n",
    "                yours.append(idx)\n",
    "                author = TEXTUAL_AVATAR\n",
    "            else:\n",
    "                author = participant_name\n",
    "\n",
    "            new_message = {\"text\": text,\n",
    "                           \"tstamp\": int(message[\"date_unixtime\"]),\n",
    "                           \"author\": author}\n",
    "            chat.append(new_message)\n",
    "            idx += 1\n",
    "\n",
    "    exported_chats[participant_name] = chat\n",
    "    your_messages_idx[participant_name] = yours\n",
    "\n",
    "\n",
    "with open('textual_avatar.csv', 'a', newline='') as csvfile:\n",
    "    writer = csv.DictWriter(csvfile, fieldnames=[\n",
    "                            \"instruction\", \"input\", \"response\"])\n",
    "    for counterpart, chat in exported_chats.items():\n",
    "        for your_reply in your_messages_idx[counterpart]:\n",
    "            last_message_to_use = max(your_reply-MAX_PREVIOUS_MESSAGES, 0)\n",
    "            your_reply_tstamp = chat[your_reply][\"tstamp\"]\n",
    "\n",
    "            for message_idx in range(last_message_to_use, your_reply):\n",
    "                if (your_reply_tstamp-chat[message_idx][\"tstamp\"]) > KEEP_CONVERSATION_SEC:\n",
    "                    last_message_to_use = message_idx+1\n",
    "                    break\n",
    "            sample = {\"counterpart\": counterpart,\n",
    "                      \"last_messages\": chat[last_message_to_use:your_reply-1],\n",
    "                      \"input\": chat[your_reply-1][\"text\"],\n",
    "                      \"text\": chat[your_reply][\"text\"]}\n",
    "            writer.writerow(convert_sample_to_instruction(\n",
    "                sample))  # adding a row\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71f3841b",
   "metadata": {},
   "source": [
    "# Add more data\n",
    "### Algorithm\n",
    "\n",
    "Download data and introduce variable path to it . Include path in final cleanup script.\n",
    "\n",
    "Detect your id, create two dictionaries with keys as person in discussion and values as messages list and indices of your replies.\n",
    "\n",
    "Iterate over saved chats resulting dictionary with ```counterpart```,```last_messages``` history list, previous ```input``` and you reply ```text```\n",
    "\n",
    "Call common ```convert_sample_to_instruction``` and add row to instruction based ```textual_avatar.csv```.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d443208",
   "metadata": {},
   "source": [
    "# Preprocess and encrypt dataset file\n",
    "\n",
    "The quality of result is highly depend on quality of training data.\n",
    "\n",
    "Your probably need to filter too long and too short answers.\n",
    "\n",
    "On other hand you may add some synthetic data using judgement from [RLHF](https://huggingface.co/blog/rlhf) or [RLAF](https://arxiv.org/pdf/2212.08073.pdf)\n",
    "\n",
    "# Encryption \n",
    "\n",
    "It is highly recommended to encrypt the file with your personal data.\n",
    "The code below store unencrypted data only in memory.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "116ad877",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install pandas --quiet\n",
    "import pandas as pd\n",
    "df = pd.read_csv('textual_avatar.csv')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6be9d307",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Encryption\n",
    "!pip install fernet --quiet\n",
    "\n",
    "from cryptography.fernet import Fernet\n",
    "import csv\n",
    "\n",
    "# Generate a random encryption key\n",
    "encryption_key = Fernet.generate_key()\n",
    "cipher_suite = Fernet(encryption_key)\n",
    "input_file = 'textual_avatar.csv'\n",
    "\n",
    "with open(input_file, 'rb') as file:\n",
    "    file_data = file.read()\n",
    "    encrypted_data = cipher_suite.encrypt(file_data)\n",
    "\n",
    "with open(input_file, 'wb') as encrypted_file:\n",
    "    encrypted_file.write(encrypted_data)\n",
    "print(\"Encryption complete. Encryption key:\", encryption_key.decode())\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70186dbb",
   "metadata": {},
   "source": [
    "### Collecting data is ending.\n",
    "### Next part you can run on another computer copying ```textual_avatar.csv` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5e9e9a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Decryption, that could be run on other machine \n",
    "!pip install fernet tqdm pandas --quiet\n",
    "import pandas as pd\n",
    "from cryptography.fernet import Fernet\n",
    "import csv\n",
    "import getpass\n",
    "import io\n",
    "\n",
    "# Ask the user for the encryption key\n",
    "encryption_key_input = getpass.getpass(prompt = 'Enter the encryption key: ')\n",
    "encryption_key = encryption_key_input.encode()\n",
    "cipher_suite = Fernet(encryption_key)\n",
    "\n",
    "input_file = 'textual_avatar.csv'\n",
    "\n",
    "with open(input_file, 'rb') as encrypted_file:\n",
    "    encrypted_data = encrypted_file.read()\n",
    "    decrypted_data = cipher_suite.decrypt(encrypted_data)\n",
    "        \n",
    "inmemorycsv=io.BytesIO(decrypted_data)\n",
    "\n",
    "df=pd.read_csv(inmemorycsv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "356f17e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "before 31205\n",
      "fitered 17368\n"
     ]
    }
   ],
   "source": [
    "print(\"before\",len(df))\n",
    "df=df[df.response.str.len()<1000]\n",
    "df=df[df.response.str.len()>2]\n",
    "print(\"fitered\",len(df))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4bf8e57",
   "metadata": {},
   "source": [
    "# Setup transformers, PEFT and LLM\n",
    "\n",
    "Here is the example of fine-tuning on LLAMA-2 model on your CUDA-enabled machine.\n",
    "For using LLAMA-2 please visit the https://ai.meta.com/resources/models-and-libraries/llama-downloads/ and accept License. \n",
    "You may also use extracted dataset for fine-tune with other models and services.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "336e11c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"transformers==4.31.0\" \"wandb\" \"scipy\" \"datasets==2.13.0\" \"peft==0.4.0\" \"accelerate==0.21.0\" \"bitsandbytes==0.40.2\" \"trl==0.4.7\" \"ninja\" \"packaging\" \"safetensors>=0.3.1\" \"cryptography\" --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6ecd7c2d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-08-25 18:58:32,177] [INFO] [real_accelerator.py:133:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "PyTorch 2.0.1\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0b8595a53a8e4bdfaee86cf8b19026a6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model, PeftModel\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments\n",
    "import torch\n",
    "print(\"PyTorch\", torch.cuda.is_available() and torch.__version__)\n",
    "model_id = \"NousResearch/Llama-2-7b-hf\"  # non-gated\n",
    "# from huggingface_hub import login\n",
    "# login()\n",
    "# model_id = \"meta-llama/Llama-2-7b-hf\" # gated\n",
    "\n",
    "\n",
    "# BitsAndBytesConfig int-4 config\n",
    "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",
    "# Load model and tokenizer\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id, quantization_config=bnb_config, use_cache=False, device_map=\"auto\")\n",
    "model.config.pretraining_tp = 1\n",
    "\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "tokenizer.padding_side = \"right\"\n",
    "\n",
    "\n",
    "# LoRA config based on QLoRA paper\n",
    "peft_config = LoraConfig(\n",
    "    lora_alpha=16,\n",
    "    lora_dropout=0.1,\n",
    "    r=64,\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")\n",
    "\n",
    "# prepare model for training\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "model = get_peft_model(model, peft_config)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a38cb64d",
   "metadata": {},
   "source": [
    "# Training script, simplified from HF tutorial\n",
    "I put bigger learning rate and loss curve become more smoother.\n",
    "It takes a time to train, on my RTX-3090 it was 8 hours. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "897acbbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datasets\n",
    "datasets.disable_caching() # don't save on disk\n",
    "datasets.config.IN_MEMORY_MAX_SIZE = 1024**3\n",
    "\n",
    "train_dataset = datasets.Dataset.from_pandas(df,preserve_index=False)\n",
    "\n",
    "\n",
    "#Alpaca format\n",
    "def format_instruction(sample):\n",
    "    return f\"\"\"### Instruction:\n",
    "{sample['instruction']}\n",
    "\n",
    "### Input:\n",
    "{sample['input']}\n",
    "\n",
    "### Response:\n",
    "{sample['response']}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "47b1d382",
   "metadata": {},
   "outputs": [],
   "source": [
    "from trl import SFTTrainer\n",
    "\n",
    "args = TrainingArguments(\n",
    "    output_dir=\"llama-7-int4-textualavatar\",\n",
    "    num_train_epochs=1,\n",
    "    per_device_train_batch_size=4,\n",
    "    gradient_accumulation_steps=2,\n",
    "    gradient_checkpointing=True,\n",
    "    optim=\"paged_adamw_32bit\",\n",
    "    logging_steps=10,\n",
    "    save_strategy=\"epoch\",\n",
    "    learning_rate=1e-3,\n",
    "    bf16=True,\n",
    "    tf32=True,\n",
    "    max_grad_norm=0.3,\n",
    "    warmup_ratio=0.03,\n",
    "    lr_scheduler_type=\"constant\",\n",
    "    disable_tqdm=True  # disable tqdm since with packing values are in correct\n",
    ")\n",
    "\n",
    "\n",
    "max_seq_length = 2048  # max sequence length for model and packing of the dataset\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    train_dataset=train_dataset,\n",
    "    peft_config=peft_config,\n",
    "    max_seq_length=max_seq_length,\n",
    "    tokenizer=tokenizer,\n",
    "    packing=True,\n",
    "    formatting_func=format_instruction,\n",
    "    args=args,\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "trainer.save_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db25dbd2",
   "metadata": {},
   "source": [
    "# Later you can load model and merge LoRA weights to 16 bit model in this script"
   ]
  },
  {
   "cell_type": "raw",
   "id": "a87640e7",
   "metadata": {},
   "source": [
    "new_model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id, torch_dtype=torch.float16, load_in_8bit=False, device_map=\"auto\", trust_remote_code=True)\n",
    "new_model = PeftModel.from_pretrained(\n",
    "    new_model, \"llama-7-int4-textualavatar/checkpoint-****\")\n",
    "model = new_model.merge_and_unload()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "918a87f0",
   "metadata": {},
   "source": [
    "# Compare with previous replies\n",
    " Run the model to chat with your previous replies and observe its responses.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c60923e8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from random import randrange\n",
    "\n",
    "\n",
    "# Load dataset and get a sample\n",
    "sample = train_dataset[randrange(len(train_dataset))]\n",
    "\n",
    "prompt = format_instruction(sample)\n",
    "print(f\"Prompt:\\n{prompt}\\n\")\n",
    "print(f\"Ground truth:\\n{sample['response']}\")\n",
    "\n",
    "input_ids = tokenizer(prompt.split(\"### Response:\")[0], return_tensors=\"pt\", truncation=True).input_ids.to(model.device)\n",
    "\n",
    "with torch.inference_mode():\n",
    "    outputs = model.generate(input_ids=input_ids, max_new_tokens=1000, do_sample=True, top_p=0.95,temperature=0.8)\n",
    "    decode = tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0]\n",
    "print(f\"Generated out :\\n{decode.split('### Response:')[1]}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "675011bd",
   "metadata": {},
   "source": [
    "# Gradio demo\n",
    "Change counterpart to the person name you want to chat."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "c8ec3380",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
      "Running on local URL:  http://127.0.0.1:7872\n",
      "\n",
      "To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7872/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": []
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "!pip install  gradio --quiet\n",
    "import gradio as gr\n",
    "\n",
    "def predict(input_text, history, counterpart=\"\"):\n",
    "    sample = {'counterpart': counterpart}\n",
    "    sample['last_messages'] = []\n",
    "    for prev_message_pair in history[-MAX_PREVIOUS_MESSAGES//2:]:\n",
    "        sample['last_messages'].append(\n",
    "            {\"author\": counterpart, \"text\": prev_message_pair[0]})\n",
    "        sample['last_messages'].append(\n",
    "            {\"author\": TEXTUAL_AVATAR, \"text\": prev_message_pair[1]})\n",
    "    sample['input'] = input_text\n",
    "    sample['text'] = \"\"\n",
    "    prompt = format_instruction(convert_sample_to_instruction(sample))\n",
    "    input_ids = tokenizer(prompt.split(\"### Response:\")[\n",
    "                          0], return_tensors=\"pt\", truncation=True).input_ids.to(model.device)\n",
    "    with torch.inference_mode():\n",
    "        outputs = model.generate(\n",
    "            input_ids=input_ids, max_new_tokens=1000, do_sample=True, top_p=0.95, temperature=0.8)\n",
    "        decode = tokenizer.batch_decode(\n",
    "            outputs.detach().cpu().numpy(), skip_special_tokens=True)[0]\n",
    "    return decode.split('### Response:')[1]\n",
    "\n",
    "gr.ChatInterface(predict).launch()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a038f587",
   "metadata": {},
   "source": [
    "# Telegram chat bot\n",
    "\n",
    "Ask @BotFather in Telegram to create a new bot, and export API_KEY as ```TELEGRAM_BOT_KEY``` system variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da6416cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install nest_asyncio python-telegram-bot --quiet\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()\n",
    "\n",
    "import telegram\n",
    "from telegram.ext import *\n",
    "\n",
    "\n",
    "import os\n",
    "\n",
    "history={}\n",
    "\n",
    "async def hello(update: telegram.Update, context: ContextTypes.DEFAULT_TYPE) -> None:\n",
    "    print.update.chat\n",
    "    await update.message.reply_text(f'Hi {update.effective_user.first_name}! I am textual avatar bot ')\n",
    "\n",
    "async def chat(update: telegram.Update, context: ContextTypes.DEFAULT_TYPE) -> None:\n",
    "    userid = f\"user{update.message.from_user.id}\"\n",
    "    if participant_names and userid in participant_names.keys(): # search in previous conversations\n",
    "        counterpart = participant_names[userid]\n",
    "    else:\n",
    "        counterpart = \"Telegram user\"\n",
    "    if counterpart not in history.keys():\n",
    "        history[counterpart]=[]\n",
    "    input_text = update.message.text\n",
    "    output = predict(input_text, history[counterpart], counterpart )\n",
    "    history[counterpart].append([input_text,output])\n",
    "    await update.message.reply_text(output)\n",
    "\n",
    "\n",
    "TELEGRAM_BOT_KEY = os.environ[\"TELEGRAM_BOT_KEY\"]\n",
    "app = ApplicationBuilder().token(TELEGRAM_BOT_KEY).build()\n",
    "app.add_handler(CommandHandler('start', hello))\n",
    "app.add_handler(MessageHandler(telegram.ext.filters.TEXT, chat))\n",
    "\n",
    "app.run_polling()"
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}