voice-cloning.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "912ea8c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchaudio\n",
    "from torch.nn.functional import softmax\n",
    "\n",
    "from IPython.display import Audio\n",
    "from scipy.io.wavfile import write as write_wav\n",
    "# from playsound import playsound\n",
    "\n",
    "from transformers import BertTokenizer\n",
    "\n",
    "import bark\n",
    "from bark import SAMPLE_RATE, generate_audio\n",
    "from bark.generation import (load_codec_model, generate_text_semantic, generate_coarse, generate_fine, codec_decode,\n",
    "    load_model)\n",
    "from encodec.utils import convert_audio\n",
    "\n",
    "def hear_fine(fine_tokens):\n",
    "    from bark.generation import SAMPLE_RATE\n",
    "    write_wav(\"/dev/shm/autoplayme.wav\", SAMPLE_RATE, codec_decode(fine_tokens))\n",
    "\n",
    "def hear_coarse(coarse_tokens):\n",
    "    hear_fine(generate_fine(coarse_tokens))\n",
    "\n",
    "def hear_semantic(semantic_tokens):\n",
    "    hear_coarse(generate_coarse(semantic_tokens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "0f81e3d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def hear_text(text, history_prompt=None, sample_count=1):\n",
    "    audio_array = np.array([], dtype=float)\n",
    "    for n in range(sample_count):\n",
    "        audio_array = np.hstack([audio_array, generate_audio(text, history_prompt=history_prompt)])\n",
    "    write_wav(f\"/dev/shm/autoplayme.wav\", SAMPLE_RATE, audio_array)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "94719971",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_history(semantic, codes, file_path):\n",
    "    if codes.shape[0] == 8: # Have fine codes\n",
    "        coarse_codes = codes[:2, :]\n",
    "        fine_codes = codes\n",
    "    elif codes.shape[0] == 2: # Have coarse codes only\n",
    "        coarse_codes = codes\n",
    "        fine_codes = generate_fine(coarse_codes)\n",
    "    else:\n",
    "        raise Exception(\"Must provide a set of coarse or fine audio tokens to save\")\n",
    "    np.savez(file_path, fine_prompt=fine_codes, coarse_prompt=coarse_codes, semantic_prompt=semantic)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "98f44f91",
   "metadata": {},
   "outputs": [],
   "source": [
    "def custom_gen_semantic(text, duration, sem_history=np.array([]), logit_mods=None, top_k=50, temp=1.0, use_gpu=True):\n",
    "    \"\"\" Custom inference method that can generate semantic tokens of a fixed duration and also return logits. \"\"\"\n",
    "    batch_size = 1 # TODO: Implement batch sizes larger than 1\n",
    "    from bark.generation import (TEXT_ENCODING_OFFSET, TEXT_PAD_TOKEN, SEMANTIC_PAD_TOKEN, SEMANTIC_INFER_TOKEN,\n",
    "        SEMANTIC_VOCAB_SIZE, SEMANTIC_RATE_HZ)\n",
    "    \n",
    "    # Sets an np array to a given length by either inserting padding or truncating\n",
    "    def set_len(arr, length, pad_value):\n",
    "        if len(arr) >= length:\n",
    "            return arr[:length]\n",
    "        else:\n",
    "            return np.pad(arr, (0, length - len(arr)), mode=\"constant\", constant_values=pad_value)\n",
    "    \n",
    "    text = bark.generation._normalize_whitespace(text)\n",
    "    assert len(text.strip()) > 0\n",
    "    \n",
    "    model_container = load_model(use_gpu=use_gpu, model_type=\"text\")\n",
    "    model = model_container[\"model\"]\n",
    "    tokenizer = model_container[\"tokenizer\"]\n",
    "    device = next(model.parameters()).device\n",
    "    \n",
    "    tokenized = np.array(bark.generation._tokenize(tokenizer, text)) + TEXT_ENCODING_OFFSET\n",
    "    x = np.hstack([set_len(tokenized, 256, TEXT_PAD_TOKEN),\n",
    "                   set_len(sem_history, 256, SEMANTIC_PAD_TOKEN), # Semantic history\n",
    "                   np.array([SEMANTIC_INFER_TOKEN])])\n",
    "    x = torch.stack([torch.from_numpy(x.astype(np.int64))] * batch_size)\n",
    "    \n",
    "    # x is a batch of 1. Its contents are 256 text tokens, then 256 semantic tokens (history), then the semantic\n",
    "    # infer token\n",
    "    prompt_len = 256 + 256 + 1\n",
    "    assert x.shape == (batch_size, prompt_len)\n",
    "    \n",
    "    kv_cache = None\n",
    "    with torch.inference_mode(), bark.generation.autocast():\n",
    "        x = x.to(device)\n",
    "        sem_toks_to_gen = max(1, math.ceil(duration * SEMANTIC_RATE_HZ))\n",
    "        assert (logit_mods is None) or (len(logit_mods) == sem_toks_to_gen)\n",
    "        logits_per_tok = []\n",
    "        for n in range(sem_toks_to_gen):\n",
    "            # print(f\"Sem token {n+1}/{sem_toks_to_gen}\")\n",
    "            \n",
    "            x_in = x[:, [-1]] if kv_cache is not None else x\n",
    "            logits, kv_cache = model(x_in, merge_context=True, use_cache=True, past_kv=kv_cache)\n",
    "            sem_logits = logits[0, 0, :SEMANTIC_VOCAB_SIZE]\n",
    "            logits_per_tok.append(sem_logits)\n",
    "            \n",
    "            if logit_mods is not None:\n",
    "                sem_logits += logit_mods[n]\n",
    "\n",
    "            _, top_indices = torch.topk(sem_logits, top_k)\n",
    "            top_k_mask = torch.full(sem_logits.shape, -math.inf).to(device)\n",
    "            top_k_mask[top_indices] = 0\n",
    "            probs = softmax((sem_logits + top_k_mask) / temp, dim=-1)\n",
    "            y = torch.multinomial(probs, num_samples=1)\n",
    "            x = torch.cat((x, y.unsqueeze(0)), dim=1)\n",
    "        \n",
    "        sem = x.detach().cpu().numpy().squeeze()[prompt_len:]\n",
    "        \n",
    "    return sem, logits_per_tok\n",
    "\n",
    "def custom_gen_coarse(semantic, use_gpu=True):\n",
    "    sliding_window_len = 60\n",
    "    \n",
    "    from bark.generation import (COARSE_RATE_HZ, SEMANTIC_RATE_HZ, N_COARSE_CODEBOOKS)\n",
    "    assert semantic.min() >= 0 and semantic.max() < SEMANTIC_VOCAB_SIZE\n",
    "    \n",
    "    # No. of coarse tokens generated per semantic token input\n",
    "    coarse_per_semantic_tok = COARSE_RATE_HZ / SEMANTIC_RATE_HZ * N_COARSE_CODEBOOKS\n",
    "    \n",
    "    # Load model\n",
    "    model = load_model(use_gpu=use_gpu, model_type=\"coarse\")\n",
    "    device = next(model.parameters()).device\n",
    "    \n",
    "    # Placeholder history data for now\n",
    "    semantic_history = np.array([], dtype=np.int32)\n",
    "    coarse_history = np.array([], dtype=np.int32)\n",
    "    \n",
    "    # Number of coarse tokens to generate\n",
    "    count_to_gen = int(np.floor(len(x) * coarse_per_semantic_tok / N_COARSE_CODEBOOKS) * N_COARSE_CODEBOOKS)\n",
    "    assert count_to_gen > 0\n",
    "    \n",
    "    count_gened = 0 # Number generated\n",
    "    \n",
    "    x_semantic = np.hstack([semantic_history, semantic]).astype(np.int32)\n",
    "    first_new_sem_index = len(semantic_history) # Index of first token of new semantic input, after history\n",
    "    x_coarse = coarse_history.astype(np.int32)\n",
    "    with torch.inference_mode(), bark.generation.autocast():\n",
    "        x_semantic = torch.from_numpy(x_semantic).unsqueeze(0).to(device)\n",
    "        x_coarse   = torch.from_numpy(x_coarse)  .unsqueeze(0).to(device)\n",
    "        count_window_steps = math.ceil(count_to_gen / sliding_window_len)\n",
    "        for _ in range(count_window_steps):\n",
    "            sem_index = first_new_sem_index + round(count_gened / coarse_per_semantic_tok)\n",
    "            \n",
    "            x_in = x_semantic[:, max(0, sem_index - max_semantic_history) :]\n",
    "            x_in = x_in[:, :256]\n",
    "        \n",
    "            logits, _ = model(x)\n",
    "\n",
    "class CoarseToSemantic:\n",
    "    def __init__(self, text, target_coarse, duration, use_gpu=True):\n",
    "        self.text = text\n",
    "        self.target_coarse = target_coarse\n",
    "        self.duration = duration\n",
    "        self.use_gpu = use_gpu\n",
    "        \n",
    "        self.logit_mod_decay = 0.99\n",
    "        self.learning_rate = 100\n",
    "        \n",
    "        self.sems = []\n",
    "        self.logit_mods = None\n",
    "        self.sim_ema = None\n",
    "        \n",
    "        self.best_sem = None\n",
    "        self.best_coarse = None\n",
    "        self.best_sem_sim = -math.inf\n",
    "\n",
    "    def run(self, step_count, silent=False):\n",
    "        for _ in range(step_count):\n",
    "            \n",
    "            # Sample semantic and coarse tokens\n",
    "            sem, logits_per_tok = custom_gen_semantic(self.text, self.duration, logit_mods=self.logit_mods, use_gpu=self.use_gpu)\n",
    "            sampled_coarse = generate_coarse(sem, use_gpu=self.use_gpu, silent=True)\n",
    "            assert sampled_coarse.shape == self.target_coarse.shape\n",
    "            \n",
    "            # Compute similarity and update EMA\n",
    "            similarity = (sampled_coarse == self.target_coarse).sum() / (sampled_coarse.shape[0] * sampled_coarse.shape[1])\n",
    "            if self.sim_ema is None:\n",
    "                self.sim_ema = similarity\n",
    "            else:\n",
    "                self.sim_ema = self.sim_ema * 0.9 + similarity * 0.1\n",
    "            if similarity > self.best_sem_sim:\n",
    "                self.best_sem = sem\n",
    "                self.best_coarse = sampled_coarse\n",
    "                self.best_sem_sim = similarity\n",
    "            \n",
    "            # Update logit mods\n",
    "            if self.logit_mods is None:\n",
    "                self.logit_mods = [torch.zeros_like(logits_per_tok[0]) for n in range(len(logits_per_tok))]\n",
    "            for sem_pos in range(len(self.logit_mods)):\n",
    "                sem_tok = sem[sem_pos]\n",
    "                self.logit_mods[sem_pos] *= self.logit_mod_decay\n",
    "                self.logit_mods[sem_pos][sem_tok] += self.learning_rate * (similarity - self.sim_ema)\n",
    "            \n",
    "            if not silent:\n",
    "                print(f\"{len(self.sems)+1}. {round(100 * similarity, 3)}%\")\n",
    "            \n",
    "            self.sems.append(sem)\n",
    "        \n",
    "        return self.sems\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37d0ff51",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecae308b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "d62bf5d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "0\n",
      "1\n",
      "1\n",
      "1\n",
      "2\n",
      "2\n",
      "2\n",
      "3\n",
      "3\n",
      "3\n",
      "4\n",
      "4\n",
      "4\n",
      "5\n"
     ]
    }
   ],
   "source": [
    "from bark.generation import (COARSE_RATE_HZ, SEMANTIC_RATE_HZ, N_COARSE_CODEBOOKS)\n",
    "coarse_per_semantic_tok = COARSE_RATE_HZ / SEMANTIC_RATE_HZ * N_COARSE_CODEBOOKS\n",
    "count_to_gen = int(np.floor(55 * coarse_per_semantic_tok / N_COARSE_CODEBOOKS) * N_COARSE_CODEBOOKS)\n",
    "for n in range(15):\n",
    "    print(round(n / coarse_per_semantic_tok))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "9751ed3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "hear_fine(en_speaker_3[\"fine_prompt\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "7e850e21",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:05<00:00, 18.29it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 19/19 [00:15<00:00,  1.21it/s]\n"
     ]
    }
   ],
   "source": [
    "hear_text(\"Do not be alarmed, this is just a test.\", \"en_speaker_3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45924478",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "564564c6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f36f704a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d444cfd",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ae95b3e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5366a73b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd6141f7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12900dd7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee3e5301",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5363b175",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:05<00:00, 16.90it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [00:17<00:00,  1.17it/s]\n"
     ]
    }
   ],
   "source": [
    "hear_text(\"Testing 1 2 3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "bad59e17",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:02<00:00,  1.38it/s]\n"
     ]
    }
   ],
   "source": [
    "sem, logits_per_tok = custom_gen_semantic(\"This is a test.\", 1.5)\n",
    "test_coarse = generate_coarse(sem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50c36140",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# c2s = CoarseToSemantic(sample_text, sample_coarse, sample_duration)\n",
    "sems = c2s.run(150)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "f415a216",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "511"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(c2s.sems)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "3c2c94be",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:02<00:00,  1.38it/s]\n"
     ]
    }
   ],
   "source": [
    "sem, _ = custom_gen_semantic(\"This is a test.\", 1.5)\n",
    "hear_semantic(sem)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c8724c6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_history(c2s.best_sem, sample_codes, \"/dev/shm/lara-test.npz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "524fc1d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:03<00:00, 31.71it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 12/12 [00:11<00:00,  1.03it/s]\n"
     ]
    }
   ],
   "source": [
    "hear_text(\"Testing 1 2 3\", history_prompt=\"/dev/shm/lara-test.npz\", sample_count=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 253,
   "id": "efb24b2f",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "only Tensors of floating point and complex dtype can require gradients",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[253], line 16\u001b[0m\n\u001b[1;32m     14\u001b[0m logits1\u001b[38;5;241m.\u001b[39mretain_grad()\n\u001b[1;32m     15\u001b[0m x2 \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mcat((x1, sample_sem(logits1)\u001b[38;5;241m.\u001b[39munsqueeze(\u001b[38;5;241m0\u001b[39m)), dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m)\n\u001b[0;32m---> 16\u001b[0m \u001b[43mx2\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrequires_grad\u001b[49m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m     17\u001b[0m x2\u001b[38;5;241m.\u001b[39mretain_grad()\n\u001b[1;32m     18\u001b[0m logits2 \u001b[38;5;241m=\u001b[39m model(x2, merge_context\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)[\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m0\u001b[39m, :SEMANTIC_VOCAB_SIZE]\n",
      "\u001b[0;31mRuntimeError\u001b[0m: only Tensors of floating point and complex dtype can require gradients"
     ]
    }
   ],
   "source": [
    "def sample_sem(logits):\n",
    "    probs = softmax(logits, dim=-1)\n",
    "    return torch.multinomial(probs, num_samples=1)\n",
    "\n",
    "# logits0 = torch.ones((SEMANTIC_VOCAB_SIZE,)) / SEMANTIC_VOCAB_SIZE\n",
    "\n",
    "x0 = x0.to(device)\n",
    "\n",
    "# x0 = ...\n",
    "logits0 = model(x0, merge_context=True)[0, 0, :SEMANTIC_VOCAB_SIZE]\n",
    "logits0.retain_grad()\n",
    "x1 = torch.cat((x0, sample_sem(logits0).unsqueeze(0)), dim=1)\n",
    "logits1 = model(x1, merge_context=True)[0, 0, :SEMANTIC_VOCAB_SIZE]\n",
    "logits1.retain_grad()\n",
    "x2 = torch.cat((x1, sample_sem(logits1).unsqueeze(0)), dim=1)\n",
    "x2.requires_grad = True\n",
    "x2.retain_grad()\n",
    "logits2 = model(x2, merge_context=True)[0, 0, :SEMANTIC_VOCAB_SIZE]\n",
    "logits2.retain_grad()\n",
    "# and so forth..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "23081cb0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:04<00:00, 20.08it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [00:16<00:00,  1.20it/s]\n"
     ]
    }
   ],
   "source": [
    "arr = generate_audio(\"What would you do with that kind of power?\", history_prompt=\"en_speaker_3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "id": "0a89d33e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:09<00:00, 10.32it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 37/37 [00:32<00:00,  1.15it/s]\n"
     ]
    }
   ],
   "source": [
    "arr2 = generate_audio(\"Do not be alarmed. This is only a test.\", history_prompt=\"en_speaker_3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "82628962",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:02<00:00, 34.13it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 11/11 [00:09<00:00,  1.19it/s]\n"
     ]
    }
   ],
   "source": [
    "arr3 = generate_audio(\"Are you hearing me? Like literally, are do you hear what I'm saying?\", history_prompt=\"en_speaker_3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "b2a983a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:04<00:00, 23.23it/s]\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 17/17 [00:14<00:00,  1.20it/s]\n"
     ]
    }
   ],
   "source": [
    "arr4 = generate_audio(\"I don't think this will work but I'll give it one last try.\", history_prompt=\"en_speaker_3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "id": "0724a48e",
   "metadata": {},
   "outputs": [],
   "source": [
    "write_wav(\"/dev/shm/autoplayme.wav\", SAMPLE_RATE, arr2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "ab6d14ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "en_speaker_3 = np.load(os.path.join(os.getcwd(), \"bark\", \"assets\", \"prompts\", \"en_speaker_3.npz\"))\n",
    "hear_fine(en_speaker_3[\"fine_prompt\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6061137",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "2ea8ca5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "codec_model = load_codec_model(use_gpu=True)\n",
    "tokenizer = BertTokenizer.from_pretrained(\"bert-base-multilingual-cased\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "ab29ebd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "audio_filepath = os.path.join(os.getcwd(), \"misc/lara-silver-box.wav\")\n",
    "sample_text = \"Dominguez said that he would use the silver box of Ix Chel to remake the world.\"\n",
    "sample_voice_name = \"nu-lara-v6\"\n",
    "\n",
    "device = \"cuda\"\n",
    "wav, sr = torchaudio.load(audio_filepath)\n",
    "wav = convert_audio(wav, sr, codec_model.sample_rate, codec_model.channels)\n",
    "wav = wav.unsqueeze(0).to(device)\n",
    "\n",
    "# Encode audio sample with EnCodec\n",
    "with torch.no_grad():\n",
    "    encoded_frames = codec_model.encode(wav)\n",
    "sample_codes = torch.cat([encoded[0] for encoded in encoded_frames], dim=-1).squeeze()  # [n_q, T]\n",
    "sample_codes = sample_codes.cpu().numpy()\n",
    "sample_coarse = sample_codes[:2, :]\n",
    "sample_duration = wav.shape[-1] / codec_model.sample_rate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0227e707",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "e677a30c",
   "metadata": {},
   "source": [
    "## First try using genetic algorithms. Doesn't work."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "508d106f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:03<00:00, 32.57it/s]\n"
     ]
    }
   ],
   "source": [
    "# generate semantic tokens\n",
    "semantic_tokens = generate_text_semantic(transcription, max_gen_duration_s=seconds, top_k=40, top_p=0.8, temp=1.0)\n",
    "# INSTEAD:\n",
    "# semantic_tokens = inverse_generate_semantic_tokens(model, coarse, transcription)\n",
    "\n",
    "output_path = \"bark/assets/prompts/\" + voice_name + \".npz\"\n",
    "np.savez(output_path, fine_prompt=codes, coarse_prompt=coarse, semantic_prompt=semantic_tokens)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffb355e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def inverse_generate_semantic_tokens(x_coarse):\n",
    "    \"\"\"Generate semantic tokens from coarse audio codes given some text.\"\"\"\n",
    "    # TODO: This is what would be needed for proper voice cloning\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "e301b670",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bark.generation import load_model, generate_coarse, generate_fine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "e6a1d4e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_container = load_model(use_gpu=True, model_type=\"text\")\n",
    "sem_model = model_container[\"model\"]\n",
    "tokenizer = model_container[\"tokenizer\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "c3b1ad0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "coarse_model = load_model(use_gpu=True, model_type=\"coarse\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 401,
   "id": "029a7cc3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def fit_coarse_and_semantic_len(text, coarse, duration):\n",
    "    \"\"\" Fits together the lengths of a text prompt, a set of coarse tokens, and an audio duration in secs.\n",
    "    Returns a length for a semantic token seq that will produce a coarse seq of the appropriate length and the\n",
    "    set of coarse tokens, possibly truncated by one step in length. \"\"\"\n",
    "    from math import ceil\n",
    "    from bark.generation import SEMANTIC_RATE_HZ, COARSE_RATE_HZ, N_COARSE_CODEBOOKS\n",
    "    semantic_to_coarse_ratio = COARSE_RATE_HZ / SEMANTIC_RATE_HZ * N_COARSE_CODEBOOKS\n",
    "    target_sem_len = ceil(duration * SEMANTIC_RATE_HZ) # Req number of semantic tokens\n",
    "    target_coarse_width = int(np.floor(target_sem_len * semantic_to_coarse_ratio / N_COARSE_CODEBOOKS)) # No. of cols in coarse translation of target_sem_len semantic tokens\n",
    "    if coarse.shape[1] > target_coarse_width:\n",
    "        coarse = coarse[:,:-1] # Drop last column\n",
    "    assert target_coarse_width == coarse.shape[1]    \n",
    "    return target_sem_len, coarse\n",
    "\n",
    "def generate_semantic_variations(text, target_coarse, duration, count, history_prompt=None, silent=False):\n",
    "    \"\"\" Randomly generates sets of semantic tokens from a text prompt with a \"target\" coarse encoding. The sets\n",
    "    will be length-matched to the target and compared with it. Returns a list of \"count\" (semantic, coarse,\n",
    "    similarity) triples. \"\"\"\n",
    "    target_sem_len, target_coarse = fit_coarse_and_semantic_len(text, target_coarse, duration)\n",
    "    tr = []\n",
    "    attempt = 1\n",
    "    while len(tr) < count:\n",
    "        if not silent:\n",
    "            print(f\"Attempt {attempt}/?; Have {len(tr)}/{count} variations\")\n",
    "        sem = generate_text_semantic(text, history_prompt, max_gen_duration_s=duration, top_k=75, top_p=None, temp=1.5, silent=True)\n",
    "        if sem.shape[0] == target_sem_len:\n",
    "            coarse = generate_coarse(sem, history_prompt, silent=True)\n",
    "            assert coarse.shape == target_coarse.shape\n",
    "            sim = (coarse == target_coarse).sum() # Similarity = sum of matching tokens\n",
    "            tr.append((sem, coarse, sim))\n",
    "        elif not silent:\n",
    "            print(f\"\\tSem length {sem.shape[0]} unusable, need {target_sem_len}\")\n",
    "        attempt += 1\n",
    "    if not silent:\n",
    "        print(\"Done\")\n",
    "    return tr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 325,
   "id": "198be2ce",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Attempt 1/?; Have 0/2 variations\n",
      "Attempt 2/?; Have 1/2 variations\n",
      "Done\n"
     ]
    }
   ],
   "source": [
    "voices = generate_semantic_variations(transcription, coarse, seconds, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "id": "41cfc491",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "350"
      ]
     },
     "execution_count": 132,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    encoded_frames = model.encode(wav)\n",
    "codes = torch.cat([encoded[0] for encoded in encoded_frames], dim=-1).squeeze()  # [n_q, T]\n",
    "\n",
    "coarse = codes[:2, :]\n",
    "\n",
    "round(seconds * SEMANTIC_RATE_HZ) # This is how many semantic tokens we want"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "id": "263d2b43",
   "metadata": {},
   "outputs": [],
   "source": [
    "coarse = coarse[:, :526]\n",
    "coarse = coarse.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 183,
   "id": "8efcc3dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "pop = [(voices[n][0], voices[n][1], sims[n]) for n in range(len(voices))]\n",
    "pop.sort(key = lambda p: -p[2]) # Sort most -> least similar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "id": "fac2e04d",
   "metadata": {},
   "outputs": [],
   "source": [
    "hear_coarse(more_voices[1][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 403,
   "id": "8169ef80",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evolve(text, target_coarse, duration, initial_pop=[], target_sim=0.5, history_prompt=None, max_gen_count=100):\n",
    "    pop_size = 50\n",
    "    new_fraction = 0.70 # Fraction of the population that is new each generation\n",
    "    \n",
    "    _, target_coarse = fit_coarse_and_semantic_len(text, target_coarse, duration)\n",
    "    coarse_count = target_coarse.shape[0] * target_coarse.shape[1]\n",
    "    \n",
    "    def get_sim_percent(voice):\n",
    "        return round(100 * voice[2] / coarse_count, 2)\n",
    "    \n",
    "    def get_avg_sim_percent(voice_list):\n",
    "        sims = [get_sim_percent(v) for v in voice_list]\n",
    "        return round(sum(sims) / len(sims), 2)\n",
    "    \n",
    "    pop = initial_pop\n",
    "    \n",
    "    # Generate initial population\n",
    "    needed_initial = pop_size - len(initial_pop)\n",
    "    if needed_initial > 0:\n",
    "        print(f\"Generating {needed_initial} voices for initial population\")\n",
    "        pop.extend(generate_semantic_variations(text, target_coarse, duration, needed_initial, history_prompt, silent=True))\n",
    "                   \n",
    "    pop.sort(key = lambda p: -p[2])\n",
    "    best = pop[0]\n",
    "    print(f\"Initial pop avg similarity: {get_avg_sim_percent(pop)}%, best similarity: {get_sim_percent(best)}%\")\n",
    "    \n",
    "    gen_count = 1\n",
    "    while (best[2] / coarse_count < target_sim) and (gen_count <= max_gen_count):\n",
    "        print(f\"Generation {gen_count}:\")\n",
    "        \n",
    "        # print(\"\\tCombining old voices...\")\n",
    "        weights = np.array([p[2] for p in pop], dtype=float)\n",
    "        weights /= weights.sum()\n",
    "        new_pop = [best]\n",
    "        while len(new_pop) < round(pop_size * (1.0 - new_fraction)):\n",
    "            # Randomly combine two voices to produce a new set of semantic tokens\n",
    "            (i_dad, i_mom) = np.random.choice(np.arange(pop_size), size=2, replace=False, p=weights)\n",
    "            dad = pop[int(i_dad)]\n",
    "            mom = pop[int(i_mom)]\n",
    "            prob = dad[2] / (dad[2] + mom[2])\n",
    "            combo = (np.random.rand(*dad[0].shape) < prob).astype(float)\n",
    "            new_sem = dad[0] * combo + mom[0] * (1.0 - combo)\n",
    "            \n",
    "            new_coarse = generate_coarse(new_sem, silent=True)\n",
    "            assert new_coarse.shape == target_coarse.shape\n",
    "            new_sim = (new_coarse == target_coarse).sum()\n",
    "            new_pop.append((new_sem, new_coarse, new_sim))\n",
    "                \n",
    "        print(f\"\\tCombinations have avg of {get_avg_sim_percent(pop[1:])}% similarity\")\n",
    "        \n",
    "        # print(\"\\tGenerating new voices...\")\n",
    "        new_voices = generate_semantic_variations(text, target_coarse, duration, pop_size - len(new_pop), history_prompt, silent=True)\n",
    "        print(f\"\\tNew voices have avg of {get_avg_sim_percent(new_voices)}% similarity\")\n",
    "        new_pop.extend(new_voices)\n",
    "        \n",
    "        pop = new_pop\n",
    "        pop.sort(key = lambda p: -p[2])\n",
    "        best = pop[0]\n",
    "        print(f\"\\tBest voice this gen has {get_sim_percent(best)}% similarity\")\n",
    "        \n",
    "        gen_count += 1\n",
    "    \n",
    "    return pop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 404,
   "id": "8dc18b81",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial pop avg similarity: 3.5%, best similarity: 13.24%\n",
      "Generation 1:\n",
      "\tCombinations have avg of 3.3% similarity\n",
      "\tNew voices have avg of 4.03% similarity\n",
      "\tBest voice this gen has 13.24% similarity\n",
      "Generation 2:\n",
      "\tCombinations have avg of 3.81% similarity\n",
      "\tNew voices have avg of 2.52% similarity\n",
      "\tBest voice this gen has 13.24% similarity\n",
      "Generation 3:\n",
      "\tCombinations have avg of 2.76% similarity\n",
      "\tNew voices have avg of 3.99% similarity\n",
      "\tBest voice this gen has 13.24% similarity\n",
      "Generation 4:\n",
      "\tCombinations have avg of 3.72% similarity\n",
      "\tNew voices have avg of 3.95% similarity\n",
      "\tBest voice this gen has 13.24% similarity\n",
      "Generation 5:\n",
      "\tCombinations have avg of 3.48% similarity\n",
      "\tNew voices have avg of 3.7% similarity\n",
      "\tBest voice this gen has 16.18% similarity\n",
      "Generation 6:\n",
      "\tCombinations have avg of 3.48% similarity\n",
      "\tNew voices have avg of 4.16% similarity\n",
      "\tBest voice this gen has 16.18% similarity\n",
      "Generation 7:\n",
      "\tCombinations have avg of 3.48% similarity\n",
      "\tNew voices have avg of 2.56% similarity\n",
      "\tBest voice this gen has 16.18% similarity\n",
      "Generation 8:\n",
      "\tCombinations have avg of 2.58% similarity\n",
      "\tNew voices have avg of 3.49% similarity\n",
      "\tBest voice this gen has 16.18% similarity\n",
      "Generation 9:\n",
      "\tCombinations have avg of 3.6% similarity\n",
      "\tNew voices have avg of 3.07% similarity\n",
      "\tBest voice this gen has 16.18% similarity\n",
      "Generation 10:\n",
      "\tCombinations have avg of 2.82% similarity\n",
      "\tNew voices have avg of 4.2% similarity\n",
      "\tBest voice this gen has 16.18% similarity\n",
      "Generation 11:\n",
      "\tCombinations have avg of 4.26% similarity\n",
      "\tNew voices have avg of 3.49% similarity\n",
      "\tBest voice this gen has 16.18% similarity\n",
      "Generation 12:\n",
      "\tCombinations have avg of 3.6% similarity\n",
      "\tNew voices have avg of 2.69% similarity\n",
      "\tBest voice this gen has 16.18% similarity\n",
      "Generation 13:\n",
      "\tCombinations have avg of 2.73% similarity\n",
      "\tNew voices have avg of 2.86% similarity\n",
      "\tBest voice this gen has 16.18% similarity\n",
      "Generation 14:\n",
      "\tCombinations have avg of 2.82% similarity\n",
      "\tNew voices have avg of 3.65% similarity\n",
      "\tBest voice this gen has 16.18% similarity\n",
      "Generation 15:\n",
      "\tCombinations have avg of 3.42% similarity\n",
      "\tNew voices have avg of 3.95% similarity\n",
      "\tBest voice this gen has 16.18% similarity\n",
      "Generation 16:\n",
      "\tCombinations have avg of 3.3% similarity\n",
      "\tNew voices have avg of 2.81% similarity\n",
      "\tBest voice this gen has 17.65% similarity\n",
      "Generation 17:\n",
      "\tCombinations have avg of 2.91% similarity\n",
      "\tNew voices have avg of 4.24% similarity\n",
      "\tBest voice this gen has 17.65% similarity\n",
      "Generation 18:\n",
      "\tCombinations have avg of 3.72% similarity\n",
      "\tNew voices have avg of 3.53% similarity\n",
      "\tBest voice this gen has 17.65% similarity\n",
      "Generation 19:\n",
      "\tCombinations have avg of 3.48% similarity\n",
      "\tNew voices have avg of 3.15% similarity\n",
      "\tBest voice this gen has 17.65% similarity\n",
      "Generation 20:\n",
      "\tCombinations have avg of 3.15% similarity\n",
      "\tNew voices have avg of 3.86% similarity\n",
      "\tBest voice this gen has 17.65% similarity\n",
      "Generation 21:\n",
      "\tCombinations have avg of 4.11% similarity\n",
      "\tNew voices have avg of 3.78% similarity\n",
      "\tBest voice this gen has 17.65% similarity\n",
      "Generation 22:\n",
      "\tCombinations have avg of 3.42% similarity\n",
      "\tNew voices have avg of 4.33% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 23:\n",
      "\tCombinations have avg of 4.47% similarity\n",
      "\tNew voices have avg of 3.53% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 24:\n",
      "\tCombinations have avg of 3.93% similarity\n",
      "\tNew voices have avg of 3.65% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 25:\n",
      "\tCombinations have avg of 3.42% similarity\n",
      "\tNew voices have avg of 3.61% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 26:\n",
      "\tCombinations have avg of 3.75% similarity\n",
      "\tNew voices have avg of 4.28% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 27:\n",
      "\tCombinations have avg of 4.08% similarity\n",
      "\tNew voices have avg of 4.03% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 28:\n",
      "\tCombinations have avg of 4.14% similarity\n",
      "\tNew voices have avg of 2.86% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 29:\n",
      "\tCombinations have avg of 3.24% similarity\n",
      "\tNew voices have avg of 3.74% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 30:\n",
      "\tCombinations have avg of 4.2% similarity\n",
      "\tNew voices have avg of 2.94% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 31:\n",
      "\tCombinations have avg of 3.39% similarity\n",
      "\tNew voices have avg of 3.28% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 32:\n",
      "\tCombinations have avg of 3.6% similarity\n",
      "\tNew voices have avg of 4.24% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 33:\n",
      "\tCombinations have avg of 4.29% similarity\n",
      "\tNew voices have avg of 3.19% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 34:\n",
      "\tCombinations have avg of 3.42% similarity\n",
      "\tNew voices have avg of 2.65% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 35:\n",
      "\tCombinations have avg of 3.21% similarity\n",
      "\tNew voices have avg of 3.15% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 36:\n",
      "\tCombinations have avg of 3.3% similarity\n",
      "\tNew voices have avg of 3.36% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 37:\n",
      "\tCombinations have avg of 3.6% similarity\n",
      "\tNew voices have avg of 3.49% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 38:\n",
      "\tCombinations have avg of 3.54% similarity\n",
      "\tNew voices have avg of 4.45% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 39:\n",
      "\tCombinations have avg of 4.26% similarity\n",
      "\tNew voices have avg of 3.4% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 40:\n",
      "\tCombinations have avg of 3.78% similarity\n",
      "\tNew voices have avg of 2.77% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 41:\n",
      "\tCombinations have avg of 2.88% similarity\n",
      "\tNew voices have avg of 3.95% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 42:\n",
      "\tCombinations have avg of 3.9% similarity\n",
      "\tNew voices have avg of 2.86% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 43:\n",
      "\tCombinations have avg of 3.39% similarity\n",
      "\tNew voices have avg of 2.81% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 44:\n",
      "\tCombinations have avg of 2.88% similarity\n",
      "\tNew voices have avg of 4.12% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 45:\n",
      "\tCombinations have avg of 4.41% similarity\n",
      "\tNew voices have avg of 3.15% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 46:\n",
      "\tCombinations have avg of 3.21% similarity\n",
      "\tNew voices have avg of 3.74% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 47:\n",
      "\tCombinations have avg of 3.57% similarity\n",
      "\tNew voices have avg of 4.07% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 48:\n",
      "\tCombinations have avg of 4.26% similarity\n",
      "\tNew voices have avg of 3.86% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 49:\n",
      "\tCombinations have avg of 3.45% similarity\n",
      "\tNew voices have avg of 3.28% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n",
      "Generation 50:\n",
      "\tCombinations have avg of 3.15% similarity\n",
      "\tNew voices have avg of 2.81% similarity\n",
      "\tBest voice this gen has 19.12% similarity\n"
     ]
    }
   ],
   "source": [
    "the_his = \"bark/assets/prompts/lara-the.npz\"\n",
    "pop = evolve(transcription, coarse, seconds, pop, target_sim=0.20, history_prompt=the_his, max_gen_count=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 344,
   "id": "25177df2",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_path = \"bark/assets/prompts/lara-the.npz\"\n",
    "np.savez(output_path, fine_prompt=codes, coarse_prompt=coarse, semantic_prompt=pop[0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 346,
   "id": "ec294d3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "the_best = pop[0]\n",
    "the_pop = pop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 433,
   "id": "91234c35",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.006012024048096"
      ]
     },
     "execution_count": 433,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "semantic_to_coarse_ratio"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bark",
   "language": "python",
   "name": "bark"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}