generate-chat-dataset.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating chat datasets\n",
    "\n",
    "This notebook is split into two parts, both of which can take a very long time to run. We start by finding the most high quality (relevant) prompts in some already high quality training data, and combine these prompts a system-prompt primed larger LLM for generating new responses that follow the ideal behavior of our fine-tuned LLM. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Finding relevant samples\n",
    "\n",
    "Load in the user chat prompts from some prominent open datasets, find the N most similar examples to a given prompt you could see in your downstream task, and save for later use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from ast import literal_eval\n",
    "import pandas as pd\n",
    "import ollama\n",
    "from tqdm import tqdm\n",
    "tqdm.pandas()\n",
    "import json\n",
    "import os \n",
    "from sklearn.metrics.pairwise import cosine_similarity\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "lmsys_ds = load_dataset(\"lmsys/lmsys-arena-human-preference-55k\")\n",
    "ultrachat_ds = load_dataset(\"HuggingFaceH4/ultrachat_200k\")\n",
    "\n",
    "lmsys_chat_prompts = [literal_eval(x['prompt'])[0] for x in lmsys_ds['train']]\n",
    "ultra_chat_prompts = [x['prompt'] for x in ultrachat_ds['train_sft']]\n",
    "\n",
    "\n",
    "lmsys_prompts_df = pd.DataFrame(lmsys_chat_prompts, columns=['prompt'])\n",
    "ultra_prompts_df = pd.DataFrame(ultra_chat_prompts, columns=['prompt'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lmsys_chat_prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def embed_prompt(prompt):\n",
    "    try:\n",
    "        return ollama.embeddings(model='nomic-embed-text:latest', prompt=prompt)\n",
    "    except:\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Please note this will take a while to run!! If you want it to run quicker, you can delete ultrachat, as that is way larger than lmsys\n",
    "# Another TODO: is to apply some sort of parallelization to this\n",
    "\n",
    "lmsys_prompts_df['embeddings'] = lmsys_prompts_df['prompt'].progress_apply(embed_prompt)\n",
    "ultra_prompts_df['embeddings'] = ultra_prompts_df['prompt'].progress_apply(embed_prompt)\n",
    "lmsys_prompts_df['embedding_arr'] = lmsys_prompts_df['embeddings'].apply(lambda x: x['embedding'] if x is not None else None)\n",
    "ultra_prompts_df['embedding_arr'] = ultra_prompts_df['embeddings'].apply(lambda x: x['embedding'] if x is not None else None)\n",
    "lmsys_prompts_df.dropna(subset=['embedding_arr'], inplace=True)\n",
    "ultra_prompts_df.dropna(subset=['embedding_arr'], inplace=True)\n",
    "\n",
    "ultra_prompts_df['embedding_shape'] = ultra_prompts_df['embedding_arr'].apply(lambda x: len(x))\n",
    "ultra_prompts_df = ultra_prompts_df[ultra_prompts_df['embedding_shape'] != 0]\n",
    "\n",
    "ultra_prompts_df.reset_index(drop=True, inplace=True)\n",
    "lmsys_prompts_df.reset_index(drop=True, inplace=True)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find most similar prompts to some seed questions. In this case, these are the kinds of questions you hope the model can answer well. \n",
    "# (Please don't judge me for some of these ;))\n",
    "\n",
    "def find_most_similar_prompts(seed_prompt, prompts_df, top_n=50):\n",
    "    seed_embedding = embed_prompt(seed_prompt)['embedding']\n",
    "    \n",
    "    similarities = cosine_similarity([seed_embedding], prompts_df['embedding_arr'].tolist())\n",
    "\n",
    "    top_indices = similarities.argsort()[0][::-1][:top_n]\n",
    "    return prompts_df.iloc[top_indices]\n",
    "\n",
    "\n",
    "\n",
    "seed_prompts = [\"Why is the sky blue\", \"Who are you?\", \"What is the meaning of life?\", \"What is the capital of France\", \"Tell me some fun facts about space\",\n",
    "                \"Favorite movie?\", \"What is the best programming language?\", \"What is the best book you have read?\", \"What is the best food?\", \"What is the best music genre?\",\n",
    "                \"How to start a fire?\", \"Any advice on how to make friends?\", \"Do you like cats?\", \"What is the square root of 144\", \"What is your name?\", \"Tell me a joke\",\n",
    "                \"Astronomy fun facts\", \"Who is the US President?\", \"Simply explain linear regression\", \"Who is George RR Martin?\", \"Who wrote the book '1984'?\", \"Is the earth flat?\",\n",
    "                \"I am bored\", \"How to make a cake\", \"How can I start a business?\", \"How to navigate in the forest\", \"How to make friends\",\n",
    "                \"What to do in NYC?\", \"What are some fun facts about the human body?\", \"When did star wars come out?\", \"Can you summarize this\", \"Let's play a game\", \"Who is better, Leo Tolstoy or Fyodor Dostoevsky?\",\"I want to go for a run, how far should I go?\",\n",
    "                \"How to make cookies\", \"Say hi to me and my friends\", \"What book should I buy next?\", \"How to prevent climate change\", \"What books did Mary Shelly write?\",\n",
    "                \"Where is Amelia Earhart?\"\n",
    "                ]\n",
    "\n",
    "lmsys_similar_prompts = {seed_prompt: find_most_similar_prompts(seed_prompt, lmsys_prompts_df) for seed_prompt in seed_prompts}\n",
    "lmsys_similar_prompts_df = pd.concat(lmsys_similar_prompts.values(), keys=lmsys_similar_prompts.keys())\n",
    "\n",
    "ultra_similar_prompts = {seed_prompt: find_most_similar_prompts(seed_prompt, ultra_prompts_df) for seed_prompt in seed_prompts}\n",
    "ultra_similar_prompts_df = pd.concat(ultra_similar_prompts.values(), keys=ultra_similar_prompts.keys())\n",
    "lmsys_similar_prompts_df = lmsys_similar_prompts_df.reset_index().rename(columns={'level_0': 'seed_prompt', 'level_1': 'idx'})\n",
    "ultra_similar_prompts_df = ultra_similar_prompts_df.reset_index().rename(columns={'level_0': 'seed_prompt', 'level_1': 'idx'})\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ultra_similar_prompts_df.sample(10)[['seed_prompt', 'prompt']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "lmsys_prompt_dataset = []\n",
    "for _, row in lmsys_similar_prompts_df.iterrows():\n",
    "    idx = row['idx']\n",
    "    lmsys_prompt_dataset.append(lmsys_ds['train'][idx]['prompt'])\n",
    "\n",
    "\n",
    "lmsys_prompt_dataset = list(set(lmsys_prompt_dataset))\n",
    "lmsys_prompt_dataset = [literal_eval(x) for x in lmsys_prompt_dataset]\n",
    "\n",
    "ultrachat_prompt_dataset = []\n",
    "for _, row in ultra_similar_prompts_df.iterrows():\n",
    "    idx = ultrachat_ds['train_sft']['prompt'].index(row['prompt'])\n",
    "    messages = ultrachat_ds['train_sft'][idx]['messages']\n",
    "    user_messages = [x['content'] for x in messages if x['role'] == 'user']\n",
    "    ultrachat_prompt_dataset.append(user_messages)\n",
    "\n",
    "\n",
    "prompt_dataset = lmsys_prompt_dataset + ultrachat_prompt_dataset + [[x] for x in seed_prompts]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save prompt dataset\n",
    "with open('data/prompt_dataset.json', 'w') as f:\n",
    "    json.dump(prompt_dataset, f, indent=4)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating the dataset\n",
    "\n",
    "Load in the relevant user prompts, craft a system prompt to match your desired model, and generate a synthetic dataset using a larger model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json \n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()\n",
    "\n",
    "from together import Together\n",
    "from groq import Groq\n",
    "\n",
    "from tqdm import tqdm\n",
    "from time import sleep\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "\n",
    "\n",
    "use_together = True\n",
    "\n",
    "if use_together:\n",
    "\n",
    "    model = 'meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo' #'meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo'\n",
    "    client = Together(api_key=os.getenv('TOGETHER_API_KEY'))\n",
    "else:\n",
    "    model = \"llama-3.1-70b-versatile\"\n",
    "    client = Groq(api_key=os.getenv('GROQ_API_KEY'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('data/prompt_dataset.json', 'r') as f:\n",
    "    prompt_dataset = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "add_everyday_conversations = True\n",
    "if add_everyday_conversations:\n",
    "    from datasets import load_dataset   \n",
    "    everyday_conversations_ds = load_dataset(\"HuggingFaceTB/everyday-conversations-llama3.1-2k\")\n",
    "    everyday_prompts = []\n",
    "    for conversation in everyday_conversations_ds['train_sft']['messages']:\n",
    "        user_messages = [x['content'] for x in conversation if x['role'] == 'user']\n",
    "        everyday_prompts.append(user_messages)    \n",
    "\n",
    "    prompt_dataset += everyday_prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "system_prompt_f = 'picard-system'\n",
    "system_prompt_path = f'prompts/{system_prompt_f}.md'\n",
    "with open(system_prompt_path, 'r') as file:\n",
    "    system_prompt = file.read()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "f = f\"{system_prompt_f}-{model.replace('/', '-')}.json\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# shuffle prompts\n",
    "import random\n",
    "random.seed(42)\n",
    "random.shuffle(prompt_dataset)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assistant_dataset = []\n",
    "\n",
    "if os.path.exists(f):\n",
    "    with open(f, 'r') as file:\n",
    "        assistant_dataset = json.load(file)\n",
    "        print(f'Loaded {len(assistant_dataset)} conversations from file')\n",
    "\n",
    "\n",
    "def create_conversation(prompts):\n",
    "    conversation = [{'role': 'system', 'content': system_prompt}]\n",
    "    try:\n",
    "        for prompt in prompts:\n",
    "            sleep(7.5)\n",
    "            conversation.append({'role': 'user', 'content': prompt})\n",
    "            response = client.chat.completions.create(\n",
    "                model=model,\n",
    "                messages=conversation,\n",
    "            )\n",
    "            conversation.append({'role': 'assistant', 'content': response.choices[0].message.content})\n",
    "        return conversation\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        return None\n",
    "\n",
    "with ThreadPoolExecutor(max_workers=3) as executor:\n",
    "    futures = [executor.submit(create_conversation, prompts) for prompts in prompt_dataset[len(assistant_dataset):]]\n",
    "    for future in tqdm(as_completed(futures), total=len(futures)):\n",
    "        result = future.result()\n",
    "        if result:\n",
    "            assistant_dataset.append(result)\n",
    "\n",
    "        # Save to file\n",
    "        with open(f, 'w') as file:\n",
    "            json.dump(assistant_dataset, file, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load assistant dataset add \"messages\" key, which is required for the dataset to be loaded by the datasets\n",
    "# TODO: remove duplicates\n",
    "with open(f, 'r') as file:\n",
    "    assistant_dataset = json.load(file)\n",
    "    assistant_dataset = [{'messages': x} for x in assistant_dataset]\n",
    "\n",
    "# Save to file\n",
    "\n",
    "f = f\"{system_prompt_f}-{model.replace('/', '-')}-messages.json\"\n",
    "\n",
    "with open(f, 'w') as file:\n",
    "    json.dump(assistant_dataset, file, indent=4)\n",
    "    \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "action-annotator",
   "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
}