public_LLarK.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Iux87x6p6U3x"
      },
      "source": [
        "## 1. Install dependencies"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9p6Im3pcNx8d"
      },
      "outputs": [],
      "source": [
        "!pip install librosa pretty_midi jams music21 datasets huggingface_hub wandb"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_lEZkuELN3cm"
      },
      "source": [
        "## Data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "72Gv5raTIPqi"
      },
      "outputs": [],
      "source": [
        "from datasets import load_dataset, concatenate_datasets\n",
        "\n",
        "urmp_dataset = load_dataset(\"jonflynn/urmp_jukebox_embeddings_qa\")\n",
        "musicnet_dataset = load_dataset(\"jonflynn/musicnet_jukebox_embeddings_abc\")\n",
        "\n",
        "urmp_train = urmp_dataset['train']\n",
        "musicnet_train = musicnet_dataset['train']\n",
        "\n",
        "combined_dataset = concatenate_datasets([urmp_train, musicnet_train])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eN7KwECklnFL"
      },
      "source": [
        "## Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vqD54ubnX_eg"
      },
      "source": [
        "### With Unsloth"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RQmNETS4YF8H"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install unsloth\n",
        "# Also get the latest nightly Unsloth!\n",
        "!pip uninstall unsloth -y && pip install --upgrade --no-cache-dir \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TzxtEPl-EUjH"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from unsloth import FastLanguageModel\n",
        "from transformers import AutoProcessor, MusicgenConfig, MusicgenForConditionalGeneration\n",
        "from transformers import AutoConfig\n",
        "from torch import Tensor\n",
        "import math\n",
        "import json\n",
        "import os\n",
        "\n",
        "class CustomLlarkModel(nn.Module):\n",
        "    def __init__(self, model_name, model_type, device, use_lora=True):\n",
        "        super(CustomLlarkModel, self).__init__()\n",
        "\n",
        "        self.model_type = model_type\n",
        "        self.device = device\n",
        "        self.target_sr = 44100\n",
        "\n",
        "        # Define special tokens for audio start and end\n",
        "        self.AUDIO_START_TOKEN = \"<AUDIO_START>\"\n",
        "        self.AUDIO_END_TOKEN = \"<AUDIO_END>\"\n",
        "\n",
        "        max_seq_length = 4096\n",
        "        dtype = torch.bfloat16  # Automatically detect dtype\n",
        "        self.language_model, self.tokenizer = FastLanguageModel.from_pretrained(\n",
        "            model_name=model_name,\n",
        "            max_seq_length=max_seq_length,\n",
        "            dtype=dtype,\n",
        "            load_in_4bit=False,\n",
        "            trust_remote_code=True\n",
        "        )\n",
        "\n",
        "        if use_lora:\n",
        "          self.language_model = FastLanguageModel.get_peft_model(\n",
        "              self.language_model,\n",
        "              r=128,\n",
        "              target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n",
        "                                \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
        "              lora_alpha=256,\n",
        "              lora_dropout=0,\n",
        "              bias=\"none\",\n",
        "              use_gradient_checkpointing=True,\n",
        "              random_state=3407\n",
        "          )\n",
        "\n",
        "        # Add special tokens to tokenizer\n",
        "        self.tokenizer.add_tokens([self.AUDIO_START_TOKEN, self.AUDIO_END_TOKEN])\n",
        "        self.language_model.resize_token_embeddings(len(self.tokenizer))\n",
        "\n",
        "        # Get token ids for special tokens\n",
        "        self.audio_start_token_id = self.tokenizer.convert_tokens_to_ids(self.AUDIO_START_TOKEN)\n",
        "        self.audio_end_token_id = self.tokenizer.convert_tokens_to_ids(self.AUDIO_END_TOKEN)\n",
        "\n",
        "        # Projection layer to map audio embeddings to language model space\n",
        "        self.audio_projection = nn.Linear(4800, self.language_model.config.hidden_size).to(device)\n",
        "\n",
        "        if self.model_type == \"llama\":\n",
        "          self.start_of_header_token = 128006\n",
        "          self.eot_token = 128001\n",
        "        elif self.model_type == \"gemma\":\n",
        "          self.start_of_turn_token = 106\n",
        "          self.eot_token = self.tokenizer.eos_token_id\n",
        "        elif self.model_type == \"qwen2\":\n",
        "          self.start_of_turn_token = 151644\n",
        "          self.eot_token = self.tokenizer.pad_token_id\n",
        "\n",
        "    def get_prompt(self, query, answer):\n",
        "        if self.model_type == \"llama\":\n",
        "            return f\"\"\"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n",
        "            You are a helpful AI assistant, you're given audio encoded as a sequence of tokens below and must transcribe it precisely.\n",
        "            {self.AUDIO_START_TOKEN}{self.AUDIO_END_TOKEN}<|eot_id|><|start_header_id|>user<|end_header_id|>\n",
        "            {query}<|eot_id|>\n",
        "            <|start_header_id|>assistant<|end_header_id|>\n",
        "            {answer}<|eot_id|><|end_of_text|>\"\"\"\n",
        "        elif self.model_type == \"gemma\":\n",
        "            return f\"\"\"<bos><start_of_turn>user\n",
        "            You're given audio encoded as a sequence of 300 tokens: {self.AUDIO_START_TOKEN}{self.AUDIO_END_TOKEN} {query}<end_of_turn>\n",
        "            <start_of_turn>model\n",
        "            {answer}<end_of_turn><eos>\"\"\"\n",
        "        elif self.model_type == \"qwen2\":\n",
        "            return f\"\"\"<|im_start|>system\n",
        "            You are a helpful AI assistant, you're given the following audio encoded as a sequence of 300 tokens and must transcribe it precisely. {self.AUDIO_START_TOKEN}{self.AUDIO_END_TOKEN}<|im_end|>\n",
        "            <|im_start|>user\n",
        "            {query}<|im_end|>\n",
        "            <|im_start|>assistant\n",
        "            {answer}<|im_end|><|endoftext|>\"\"\"\n",
        "\n",
        "    def get_query_prompt(self, query):\n",
        "        if self.model_type == \"llama\":\n",
        "            return f\"\"\"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n",
        "            You are a helpful AI assistant, you're given audio encoded as a sequence of tokens below and must transcribe it precisely.\n",
        "            {self.AUDIO_START_TOKEN}{self.AUDIO_END_TOKEN}<|eot_id|><|start_header_id|>user<|end_header_id|>\n",
        "            {query}<|eot_id|>\n",
        "            <|start_header_id|>assistant<|end_header_id|>\"\"\"\n",
        "        elif self.model_type == \"gemma\":\n",
        "            return f\"\"\"<bos><start_of_turn>user\n",
        "            You're given audio encoded as a sequence of tokens: {self.AUDIO_START_TOKEN}{self.AUDIO_END_TOKEN} {query}<end_of_turn>\n",
        "            <start_of_turn>model\"\"\"\n",
        "        elif self.model_type == \"qwen2\":\n",
        "            return f\"\"\"<|im_start|>system\n",
        "            You are a helpful AI assistant, you're given the following audio encoded as a sequence of tokens and must transcribe it precisely. {self.AUDIO_START_TOKEN}{self.AUDIO_END_TOKEN}<|im_end|>\n",
        "            <|im_start|>user\n",
        "            {query}<|im_end|>\n",
        "            <|im_start|>assistant\"\"\"\n",
        "\n",
        "    def forward(self, **kwargs):\n",
        "        audio_embedding = kwargs.get('embedding')\n",
        "        queries = kwargs.get('query')\n",
        "        answers = kwargs.get('answer')\n",
        "\n",
        "        audio_embedding = torch.tensor(audio_embedding, dtype=torch.bfloat16).to(self.device)\n",
        "\n",
        "        # Project audio embeddings to match Llama's hidden size\n",
        "        audio_features = self.audio_projection(audio_embedding).to(self.device)\n",
        "\n",
        "        prompts = [self.get_prompt(query, answer) for query, answer in zip(queries, answers)]\n",
        "\n",
        "        self.tokenizer.pad_token = self.tokenizer.eos_token\n",
        "        self.tokenizer.padding_side = \"right\"\n",
        "\n",
        "        # Tokenize the input without padding\n",
        "        tokenizer_output = self.tokenizer(\n",
        "            prompts,\n",
        "            return_tensors='pt',\n",
        "            truncation=True,\n",
        "            padding=True,\n",
        "            # max_length=max_length_without_audio,\n",
        "            # pad_to_multiple_of=max_length_without_audio\n",
        "        )\n",
        "\n",
        "        input_ids = tokenizer_output['input_ids'].to(self.device)\n",
        "        attention_mask = tokenizer_output['attention_mask'].to(self.device)\n",
        "\n",
        "        inputs_embeds = self.language_model.get_input_embeddings()(input_ids)\n",
        "        number_of_audio_tokens = audio_embedding.shape[1]\n",
        "\n",
        "        batch_size, sequence_length = input_ids.shape\n",
        "        max_seq_length = sequence_length + number_of_audio_tokens\n",
        "\n",
        "        new_inputs_embeds = torch.zeros(\n",
        "            (batch_size, max_seq_length, inputs_embeds.size(2)),\n",
        "            device=inputs_embeds.device\n",
        "        )\n",
        "        new_attention_mask = torch.zeros((batch_size, max_seq_length), device=inputs_embeds.device)\n",
        "\n",
        "        # Insert the projected audio embeddings between the AUDIO_START_TOKEN and AUDIO_END_TOKEN for each sample in the batch\n",
        "        for i in range(batch_size):\n",
        "            start_pos = (input_ids[i] == self.audio_start_token_id).nonzero(as_tuple=True)[0]\n",
        "            end_pos = (input_ids[i] == self.audio_end_token_id).nonzero(as_tuple=True)[0]\n",
        "\n",
        "            if not len(start_pos) == 1:\n",
        "                raise ValueError(f\"Incorrect number of audio start tokens in the input. Got {len(start_pos)} tokens\")\n",
        "\n",
        "            if not len(end_pos) == 1:\n",
        "                raise ValueError(f\"Incorrect number of audio end tokens in the input. Got {len(end_pos)} tokens\")\n",
        "\n",
        "            if start_pos.size(0) > 0 and end_pos.size(0) > 0:\n",
        "                start_pos = start_pos[0].item()\n",
        "                end_pos = end_pos[0].item()\n",
        "\n",
        "                # Create the new embedding sequence\n",
        "                part1 = inputs_embeds[i, :start_pos + 1]\n",
        "                part2 = audio_features[i]\n",
        "                part3 = inputs_embeds[i, end_pos:]\n",
        "\n",
        "                new_embed = torch.cat((part1, part2, part3), dim=0)\n",
        "\n",
        "                new_inputs_embeds[i] = new_embed\n",
        "\n",
        "                # Adjust attention mask for the inserted audio embeddings\n",
        "                new_attention_mask[i] = torch.cat((\n",
        "                    attention_mask[i, :start_pos + 1],\n",
        "                    torch.ones(number_of_audio_tokens, device=inputs_embeds.device),\n",
        "                    attention_mask[i, end_pos:]\n",
        "                ), dim=0)\n",
        "\n",
        "        position_ids = (new_attention_mask.cumsum(-1) - 1).masked_fill_((new_attention_mask == 0), 1).long()\n",
        "\n",
        "        labels_list = []\n",
        "\n",
        "        for i in range(input_ids.size(0)):\n",
        "            # Create a copy of input_ids to serve as the base for labels\n",
        "            sample_labels = torch.full_like(input_ids[i], -100)\n",
        "\n",
        "            if self.model_type == \"llama\":\n",
        "                # Get the third start_of_header_token which is the one where the assistant's response starts\n",
        "                assistant_start_pos = (input_ids[i] == self.start_of_header_token).nonzero(as_tuple=True)[0][2].item()\n",
        "            elif self.model_type == \"gemma\":\n",
        "                # Get the second start_of_turn_token which is the one where the model's response starts then +1\n",
        "                assistant_start_pos = (input_ids[i] == self.start_of_turn_token).nonzero(as_tuple=True)[0][1].item() + 1\n",
        "            elif self.model_type == \"qwen2\":\n",
        "                # Get the third start_of_turn_token which is the one where the model's response starts then +1\n",
        "                assistant_start_pos = (input_ids[i] == self.start_of_turn_token).nonzero(as_tuple=True)[0][2].item() + 1\n",
        "\n",
        "            eot_pos = (input_ids[i] == self.eot_token).nonzero(as_tuple=True)[0][0].item()\n",
        "\n",
        "            # Fill in the input_ids for the assistant response part\n",
        "            sample_labels[assistant_start_pos:eot_pos] = input_ids[i, assistant_start_pos:eot_pos]\n",
        "\n",
        "            # Find AUDIO_START_TOKEN and AUDIO_END_TOKEN positions\n",
        "            audio_start_pos = (input_ids[i] == self.audio_start_token_id).nonzero(as_tuple=True)[0].item()\n",
        "            audio_end_pos = (input_ids[i] == self.audio_end_token_id).nonzero(as_tuple=True)[0].item()\n",
        "\n",
        "            # Insert padding values between AUDIO_START_TOKEN and AUDIO_END_TOKEN for each sample separately\n",
        "            sample_labels = torch.cat((\n",
        "                sample_labels[:audio_start_pos + 1],\n",
        "                torch.full((number_of_audio_tokens,), -100, dtype=input_ids.dtype, device=input_ids.device),\n",
        "                sample_labels[audio_start_pos + 1:audio_end_pos + 1],\n",
        "                sample_labels[audio_end_pos + 1:]\n",
        "            ), dim=0)\n",
        "\n",
        "            # Verify the length after concatenation\n",
        "            expected_length = input_ids.size(1) + number_of_audio_tokens\n",
        "            if sample_labels.size(0) != expected_length:\n",
        "                raise ValueError(f\"Concatenation error: expected length {expected_length} but got {sample_labels.size(0)}\")\n",
        "\n",
        "            labels_list.append(sample_labels)\n",
        "\n",
        "        labels = torch.stack(labels_list)\n",
        "\n",
        "        print(f\"Sequence shape: {new_inputs_embeds.shape}\")\n",
        "        print(f\"Number of tokens used as labels: {(labels != -100).sum().item()}\")\n",
        "\n",
        "        assert max_seq_length == position_ids.size(1) == new_inputs_embeds.size(1) == labels.size(1) == new_attention_mask.size(1), \"position_ids, new_inputs_embeds, new_labels and new_attention_mask must have the same sequence length equal to the max_seq_length\"\n",
        "\n",
        "        outputs = self.language_model(\n",
        "            inputs_embeds=new_inputs_embeds,\n",
        "            position_ids=position_ids,\n",
        "            attention_mask=new_attention_mask,\n",
        "            labels=labels,\n",
        "            return_dict=True\n",
        "        )\n",
        "\n",
        "        if 'loss' in outputs:\n",
        "            print(f\"Loss: {outputs.loss.item()}\")\n",
        "\n",
        "        return outputs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NtgVobWTViL-"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eo_lV81dWWit"
      },
      "source": [
        "### Load model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JE6auga51XWJ"
      },
      "outputs": [],
      "source": [
        "from huggingface_hub import login\n",
        "\n",
        "# Log in to Hugging Face\n",
        "login(token='hf_token')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J66KMZI8exMD"
      },
      "outputs": [],
      "source": [
        "# model = CustomLlarkModel(\"unsloth/gemma-2b-it\", \"gemma\", \"cuda\", use_lora=True)\n",
        "model = CustomLlarkModel(\"Qwen/Qwen2-7B-Instruct\", \"qwen2\", \"cuda\", use_lora=True)\n",
        "# model = CustomLlarkModel(\"unsloth/Llama-3.1-Storm-8B\", \"llama\", \"cuda\", use_lora=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RF5g4AS6GtVN"
      },
      "source": [
        "### Load wandb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kMs2GDzA1fWe"
      },
      "outputs": [],
      "source": [
        "import wandb, os\n",
        "wandb.login()\n",
        "\n",
        "wandb_project = \"llark\"\n",
        "if len(wandb_project) > 0:\n",
        "    os.environ[\"WANDB_PROJECT\"] = wandb_project"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "26AZA7qId07Z"
      },
      "outputs": [],
      "source": [
        "project = \"LLarK\"\n",
        "run_name = \"run\"\n",
        "project_and_run_name = project + \"-\" + run_name\n",
        "output_dir = \"./\" + project_and_run_name"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tCMWxBCd9DIL"
      },
      "source": [
        "### Train with SFTtrainer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "McA8lfUe5b7H"
      },
      "outputs": [],
      "source": [
        "from transformers import Trainer, TrainingArguments\n",
        "from datetime import datetime\n",
        "\n",
        "wandbname = project + \"-\" + run_name\n",
        "\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "\n",
        "def custom_collate_fn(batch):\n",
        "    queries = [item['question'] for item in batch]\n",
        "    answers = [item['answer'] for item in batch]\n",
        "    embeddings = [item['embedding'] for item in batch]\n",
        "\n",
        "    # can't return tensors here or unsloth has weird error\n",
        "    return {\n",
        "        'query': queries,\n",
        "        'answer': answers,\n",
        "        'embedding': embeddings\n",
        "    }\n",
        "\n",
        "training_args = TrainingArguments(\n",
        "    output_dir=output_dir,\n",
        "    num_train_epochs=1,\n",
        "    per_device_train_batch_size = 2,\n",
        "    gradient_accumulation_steps = 1,\n",
        "    warmup_ratio = 0.1,\n",
        "    logging_dir='./logs',\n",
        "    learning_rate = 2e-5,\n",
        "    logging_steps = 1,\n",
        "    #eval_strategy=\"epoch\",\n",
        "    #eval_steps=75,\n",
        "    #save_steps=100,\n",
        "    max_grad_norm=10.0,\n",
        "    fp16 = not torch.cuda.is_bf16_supported(),\n",
        "    bf16 = torch.cuda.is_bf16_supported(),\n",
        "    optim = \"adamw_8bit\",\n",
        "    weight_decay = 0.001,\n",
        "    seed = 3407,\n",
        "    save_strategy=\"no\",\n",
        "    #save_strategy=\"epoch\",\n",
        "    lr_scheduler_type = \"cosine\",\n",
        "    #load_best_model_at_end=True,\n",
        "    remove_unused_columns=False,\n",
        "    report_to=\"wandb\",\n",
        "    run_name=f\"{wandbname}-{datetime.now().strftime('%m-%d-%H-%M')}\"\n",
        ")\n",
        "\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    train_dataset=combined_dataset,\n",
        "    data_collator=custom_collate_fn\n",
        ")\n",
        "\n",
        "trainer.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pgIQImZND3HH"
      },
      "source": [
        "## Save model and checkpoints"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AeQybjoQD18G"
      },
      "outputs": [],
      "source": [
        "torch.save(model.audio_projection.state_dict(), \"/content/llark_multi_modal_projector_weights.pth\")\n",
        "model.language_model.save_pretrained_merged(\"model_16bit_merged\", model.tokenizer, save_method = \"merged_16bit\",)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!cp -r \"/content/llark_multi_modal_projector_weights.pth\" \"/content/drive/My Drive/automatic-music-transcription/saved_models/\"\n",
        "!cp -r \"/content/model_16bit_merged\" \"/content/drive/My Drive/automatic-music-transcription/saved_models/\""
      ],
      "metadata": {
        "id": "vGwJcOvPwutI"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4EbkJr4Z9EOg"
      },
      "source": [
        "## Try model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qNw-UeDJzYMr"
      },
      "source": [
        "### Without Unsloth for inference"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Unsloth doesn't support passing in `inputs_embeds` to the `generate` function which we need to do to accommodate the audio tokens so instead we use just `transformers` for inference\n",
        "\n",
        "https://github.com/unslothai/unsloth/issues/862"
      ],
      "metadata": {
        "id": "Fj3LQe9eKudq"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Load in full 32bit otherwise there's errors"
      ],
      "metadata": {
        "id": "2v27etzC60mP"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ONVfZOzNzZWC"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from transformers import AutoProcessor, AutoTokenizer, AutoModelForCausalLM, AutoConfig, MusicgenConfig, MusicgenForConditionalGeneration\n",
        "from torch import Tensor\n",
        "import math\n",
        "from peft import PeftModel\n",
        "\n",
        "class CustomLlarkModel(nn.Module):\n",
        "    def __init__(self, language_model_name, model_type, device, use_lora=True):\n",
        "        super(CustomLlarkModel, self).__init__()\n",
        "\n",
        "        self.model_type = model_type\n",
        "        self.device = device\n",
        "\n",
        "        # Initialize tokenizer and model\n",
        "        self.tokenizer = AutoTokenizer.from_pretrained(language_model_name)\n",
        "        self.language_model = AutoModelForCausalLM.from_pretrained(language_model_name, device_map=\"auto\").to(self.device)\n",
        "        self.language_model.gradient_checkpointing_disable()\n",
        "\n",
        "        # Define special tokens for audio start and end\n",
        "        self.AUDIO_START_TOKEN = \"<AUDIO_START>\"\n",
        "        self.AUDIO_END_TOKEN = \"<AUDIO_END>\"\n",
        "\n",
        "        self.audio_start_token_id = self.tokenizer.convert_tokens_to_ids(self.AUDIO_START_TOKEN)\n",
        "        self.audio_end_token_id = self.tokenizer.convert_tokens_to_ids(self.AUDIO_END_TOKEN)\n",
        "\n",
        "        self.audio_projection = nn.Linear(4800, self.language_model.config.hidden_size).to(device)\n",
        "\n",
        "        if self.model_type == \"llama\":\n",
        "          self.start_of_header_token = 128006\n",
        "          self.eot_token = 128001\n",
        "        elif self.model_type == \"gemma\":\n",
        "          self.start_of_turn_token = 106\n",
        "          self.eot_token = self.tokenizer.eos_token_id # 1\n",
        "        elif self.model_type == \"qwen2\":\n",
        "          self.start_of_turn_token = 151644\n",
        "          self.eot_token = self.tokenizer.pad_token_id # 1\n",
        "\n",
        "    def get_prompt(self, query):\n",
        "        if self.model_type == \"llama\":\n",
        "            return f\"\"\"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n",
        "            You are a helpful AI assistant, you're given audio encoded as a sequence of 300 tokens below and must transcribe it precisely.\n",
        "            {self.AUDIO_START_TOKEN}{self.AUDIO_END_TOKEN}<|eot_id|><|start_header_id|>user<|end_header_id|>\n",
        "            {query}<|eot_id|>\n",
        "            <|start_header_id|>assistant<|end_header_id|>\"\"\"\n",
        "        elif self.model_type == \"gemma\":\n",
        "            return f\"\"\"<bos><start_of_turn>user\n",
        "            You're given audio encoded as a sequence of 300 tokens: {self.AUDIO_START_TOKEN}{self.AUDIO_END_TOKEN} {query}<end_of_turn>\n",
        "            <start_of_turn>model\"\"\"\n",
        "        elif self.model_type == \"qwen2\":\n",
        "            return f\"\"\"<|im_start|>system\n",
        "            You are a helpful AI assistant, you're given the following audio encoded as a sequence of 300 tokens and must transcribe it precisely. {self.AUDIO_START_TOKEN}{self.AUDIO_END_TOKEN}<|im_end|>\n",
        "            <|im_start|>user\n",
        "            {query}<|im_end|>\n",
        "            <|im_start|>assistant\"\"\"\n",
        "\n",
        "    def load_trained_weights(self, projector_weights_path):\n",
        "        projector_state_dict = torch.load(projector_weights_path)\n",
        "        self.audio_projection.load_state_dict(projector_state_dict)\n",
        "\n",
        "        print(\"Trained weights loaded successfully.\")\n",
        "\n",
        "    def generate(self, audio_embedding, query, max_new_tokens=4096, num_beams=1, do_sample=False, top_k=None, top_p=None, temperature=1.0):\n",
        "        audio_embedding = torch.tensor([audio_embedding], dtype=torch.bfloat16).to(self.device)\n",
        "        audio_features = self.audio_projection(audio_embedding).to(self.device)\n",
        "        number_of_audio_tokens = audio_features.shape[1]\n",
        "\n",
        "        self.tokenizer.pad_token = self.tokenizer.eos_token\n",
        "        self.tokenizer.padding_side = \"right\"\n",
        "\n",
        "        prompt = self.get_prompt(query)\n",
        "\n",
        "        # Tokenize the input without padding\n",
        "        tokenizer_output = self.tokenizer(\n",
        "            prompt,\n",
        "            return_tensors='pt',\n",
        "            truncation=True,\n",
        "        )\n",
        "\n",
        "        input_ids = tokenizer_output['input_ids'].to(self.device)\n",
        "        attention_mask = tokenizer_output['attention_mask'].to(self.device)\n",
        "\n",
        "        inputs_embeds = self.language_model.get_input_embeddings()(input_ids)\n",
        "\n",
        "        batch_size, sequence_length = input_ids.shape\n",
        "        max_seq_length = sequence_length + number_of_audio_tokens\n",
        "\n",
        "        new_inputs_embeds = torch.zeros(\n",
        "            (batch_size, max_seq_length, inputs_embeds.size(2)),\n",
        "            device=inputs_embeds.device\n",
        "        )\n",
        "        new_attention_mask = torch.zeros((batch_size, max_seq_length), device=inputs_embeds.device)\n",
        "\n",
        "        # Insert the projected audio embeddings between the AUDIO_START_TOKEN and AUDIO_END_TOKEN for each sample in the batch\n",
        "        for i in range(batch_size):\n",
        "            start_pos = (input_ids[i] == self.audio_start_token_id).nonzero(as_tuple=True)[0]\n",
        "            end_pos = (input_ids[i] == self.audio_end_token_id).nonzero(as_tuple=True)[0]\n",
        "\n",
        "            if not len(start_pos) == 1:\n",
        "                raise ValueError(f\"Incorrect number of audio start tokens in the input. Got {len(start_pos)} tokens\")\n",
        "\n",
        "            if not len(end_pos) == 1:\n",
        "                raise ValueError(f\"Incorrect number of audio end tokens in the input. Got {len(end_pos)} tokens\")\n",
        "\n",
        "            if start_pos.size(0) > 0 and end_pos.size(0) > 0:\n",
        "                start_pos = start_pos[0].item()\n",
        "                end_pos = end_pos[0].item()\n",
        "\n",
        "                # Create the new embedding sequence\n",
        "                part1 = inputs_embeds[i, :start_pos + 1]\n",
        "                part2 = audio_features[i]\n",
        "                part3 = inputs_embeds[i, end_pos:]\n",
        "\n",
        "                new_embed = torch.cat((part1, part2, part3), dim=0)\n",
        "\n",
        "                new_inputs_embeds[i] = new_embed\n",
        "\n",
        "                # Adjust attention mask for the inserted audio embeddings\n",
        "                new_attention_mask[i] = torch.cat((\n",
        "                    attention_mask[i, :start_pos + 1],\n",
        "                    torch.ones(number_of_audio_tokens, device=inputs_embeds.device),\n",
        "                    attention_mask[i, end_pos:]\n",
        "                ), dim=0)\n",
        "\n",
        "        position_ids = (new_attention_mask.cumsum(-1) - 1).masked_fill_((new_attention_mask == 0), 1).long()\n",
        "\n",
        "        # Ensure the `inputs_embeds` are used in the first step of generation\n",
        "        generation_params = {\n",
        "            \"inputs_embeds\": new_inputs_embeds,\n",
        "            \"attention_mask\": new_attention_mask,\n",
        "            \"max_new_tokens\": max_new_tokens,\n",
        "            # \"position_ids\": position_ids,    generate() currently breaks if passed in `position_ids`\n",
        "            \"num_beams\": num_beams,\n",
        "            \"use_cache\": False,\n",
        "            \"do_sample\": do_sample,\n",
        "            \"temperature\": temperature,\n",
        "        }\n",
        "\n",
        "        if top_k is not None:\n",
        "            generation_params[\"top_k\"] = top_k\n",
        "        if top_p is not None:\n",
        "            generation_params[\"top_p\"] = top_p\n",
        "\n",
        "        output_ids = self.language_model.generate(**generation_params)\n",
        "        return self.tokenizer.decode(output_ids[0], skip_special_tokens=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PVSlIk74YKNy"
      },
      "outputs": [],
      "source": [
        "!cp -r \"/content/drive/My Drive/automatic-music-transcription/saved_models/model_16bit_merged/\" \"/content/\"\n",
        "!cp -r \"/content/drive/My Drive/automatic-music-transcription/saved_models/llark_multi_modal_projector_weights.pth\" \"/content/llark_multi_modal_projector_weights.pth\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eCRCkj1KFKmi"
      },
      "outputs": [],
      "source": [
        "model = CustomLlarkModel(\"/content/model_16bit_merged\", \"qwen2\", \"cuda\", use_lora=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lnSB5FWs0fno"
      },
      "outputs": [],
      "source": [
        "model.load_trained_weights(projector_weights_path=\"/content/llark_multi_modal_projector_weights.pth\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xIHfldxEjg5u"
      },
      "outputs": [],
      "source": [
        "audio_embedding = combined_dataset[0]['embedding']\n",
        "query = combined_dataset[0]['question']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NZ6wY11V6NG9"
      },
      "outputs": [],
      "source": [
        "output_ids = model.generate(audio_embedding, query, max_new_tokens=2048, do_sample=True, top_p=0.7, temperature=1.2)\n",
        "output_ids"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "L4",
      "machine_shape": "hm",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}