{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fed7c348",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"6\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9c62b25",
   "metadata": {},
   "source": [
    "# BCT Inference\n",
    "\n",
    "This notebook is a hackable place to test BCT model inference with ASR."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23f330f5",
   "metadata": {},
   "source": [
    "## Load model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c7d0311",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.gpt.generation import load_model, GPT\n",
    "from suno_utils.utils.s3 import download_s3_file_if_needed\n",
    "import torch\n",
    "\n",
    "N_BATCH = 2\n",
    "\n",
    "gpt_model_path = \"/app2/suno/checkpoints/2025-09-19_20-19-51/last_ckpt_infer.pt\" # non-causal\n",
    "tokenizer_path = \"s3://suno-data/georg/models/tokenizers/tokenizer_60k.json\"\n",
    "\n",
    "\n",
    "model_container = load_model(\n",
    "    ckpt_path=download_s3_file_if_needed(gpt_model_path),\n",
    "    tokenizer_path=download_s3_file_if_needed(tokenizer_path),\n",
    ")\n",
    "model: GPT = model_container[\"model\"]\n",
    "assert isinstance(model, GPT)\n",
    "\n",
    "\n",
    "cfg = model.config\n",
    "# cfg.use_asr = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a1d5952",
   "metadata": {},
   "source": [
    "## Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54e8497e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from suno_utils.audio import Audio\n",
    "from suno_utils.diffusion.generation import encode_semantic\n",
    "from suno_utils.tasks.mert_25 import preload_models as preload_semantic_models_\n",
    "\n",
    "cluster_model = preload_semantic_models_(\n",
    "    checkpoint_filepath=\"s3://suno-data/georg/models/semantic/mert_25.pt\",\n",
    "    centroids_filepath=\"s3://suno-data/georg/models/semantic/mert_25_2x4k.npy\",\n",
    ")[\"cluster_model\"].cpu()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ade389df",
   "metadata": {},
   "outputs": [],
   "source": [
    "# audio_path = \"/app2/suno/data/raw_audio_opus_v0/zW5TueAUdG4.opus\"\n",
    "# audio = Audio.from_file(audio_path)\n",
    "# audio = audio.get_segment(from_s=10, to_s=30)\n",
    "\n",
    "# audio_path = \"/app2/suno/data/raw_audio_opus_v0/G3Feih2sPOI.opus\"\n",
    "# audio = Audio.from_file(audio_path)\n",
    "# audio = audio.get_segment(from_s=15, to_s=135)\n",
    "\n",
    "audio_path = \"/app2/suno/data/raw_audio_opus_v0/kfXphP18jPc.opus\" # french\n",
    "audio = Audio.from_file(audio_path)\n",
    "audio = audio.get_segment(from_s=0, to_s=30)\n",
    "\n",
    "# audio_path = \"/app2/suno/data/podcast/audio_chunks/chunk_8_6b076c76-4e61-46aa-8bc1-8d1f95afcea7.mp3.mp3\"\n",
    "# audio_path = \"/app2/suno/data/podcast/audio_chunks/chunk_16_4eb5c269-dcf6-41d2-833d-87dc8f6fafd3.mp3.mp3\"\n",
    "# audio_path = \"/app2/suno/data/raw_audio_opus_v0/Uon2ZabBRQM.opus\"\n",
    "# audio = Audio.from_file(audio_path)\n",
    "audio.play()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2e7e428",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from suno_utils.gpt.bct.bct_generation_simple import BCTGenerationConfig, BlockSequence, generate_block\n",
    "from suno_utils.gpt.bct.bct import Block, BlockType\n",
    "\n",
    "TextBlockType = BlockType(\n",
    "    name=\"text\",\n",
    "    is_causal=True,\n",
    ")\n",
    "NonCausalSemanticBlockType = BlockType(\n",
    "    name=\"continuous_semantic\",\n",
    "    is_causal=False,\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f4f8960",
   "metadata": {},
   "source": [
    "## ASR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b265fbe1",
   "metadata": {},
   "outputs": [],
   "source": [
    "codes = torch.from_numpy(encode_semantic(\n",
    "    audio, pad_to_chunksize=True, batch_size=48, do_clustering=False\n",
    ").T).unsqueeze(0) # (dim, seq_len)\n",
    "\n",
    "codes.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec2d84ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"Number of model parameters: {num_params}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6ba81da",
   "metadata": {},
   "outputs": [],
   "source": [
    "sem_block = Block(\n",
    "    NonCausalSemanticBlockType,\n",
    "    inputs={\n",
    "        \"continuous_semantic_input\": codes.to(torch.bfloat16),\n",
    "    },\n",
    ")\n",
    "prompt_text = \"[has tags]\"\n",
    "text_codes = [cfg.text_infer_token] + model_container[\"tokenizer\"].encode(prompt_text)\n",
    "print(text_codes)\n",
    "text_block = Block(\n",
    "    TextBlockType,\n",
    "    inputs={\n",
    "        \"text_input\": torch.tensor(text_codes).reshape(1, 1, -1),\n",
    "    },\n",
    ")\n",
    "\n",
    "blocks = BlockSequence([sem_block, text_block])\n",
    "no_sem_blocks = BlockSequence([text_block])\n",
    "gconf = BCTGenerationConfig([(1, blocks), (-0, no_sem_blocks)], max_autoregressive_steps=1000, eos_token=model.config.text_pad_token, top_k=1)\n",
    "\n",
    "block = generate_block(model, gconf)\n",
    "text_codes = block.inputs[\"text_input\"][0, 0, : len(block)]\n",
    "print(text_codes.shape)\n",
    "tok = model_container[\"tokenizer\"]\n",
    "\n",
    "print(tok.decode(text_codes, clean_up_tokenization_spaces=True))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e937e93a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# To see the first text logits for generation with different tag prefixes\n",
    "import torch\n",
    "from suno_utils.gpt.bct.bct_generation_simple import BCTGenerationConfig, BlockSequence\n",
    "from suno_utils.gpt.bct.bct import Block, BlockType\n",
    "from copy import deepcopy\n",
    "\n",
    "def test_prompt_with_prefix(model, model_container, codes, cfg, base_prompt_text, prefix):\n",
    "    \"\"\"Test a prompt with a specific prefix and return the logits analysis\"\"\"\n",
    "    \n",
    "    # Modify the prompt text with the prefix\n",
    "    modified_prompt_text = f\"{prefix}{base_prompt_text}\"\n",
    "    print(f\"\\n{'='*50}\")\n",
    "    print(f\"Testing with prefix: '{prefix}'\")\n",
    "    print(f\"Full prompt: '{modified_prompt_text}'\")\n",
    "    print('='*50)\n",
    "    \n",
    "    # Setup blocks\n",
    "    sem_block = Block(\n",
    "        NonCausalSemanticBlockType,\n",
    "        inputs={\n",
    "            \"continuous_semantic_input\": codes.to(torch.bfloat16),\n",
    "        },\n",
    "    )\n",
    "    \n",
    "    text_codes = [cfg.text_infer_token] + model_container[\"tokenizer\"].encode(modified_prompt_text)\n",
    "    text_block = Block(\n",
    "        TextBlockType,\n",
    "        inputs={\n",
    "            \"text_input\": torch.tensor(text_codes).reshape(1, 1, -1),\n",
    "        },\n",
    "    )\n",
    "\n",
    "    blocks = BlockSequence([sem_block, text_block])\n",
    "    gconf = BCTGenerationConfig([(1.0, blocks)], max_autoregressive_steps=1, temperature=0.9)\n",
    "\n",
    "    # Setup streams manually to capture logits\n",
    "    weight, sequence = gconf.prompts[0]\n",
    "    num_blocks = cfg.block_size // 256 * cfg.n_layer\n",
    "    kv_cache = model.setup_caches(block_size=256, num_blocks=num_blocks)\n",
    "    block_table = torch.arange(num_blocks, device=model.device,\n",
    "    dtype=torch.int32).reshape(1, cfg.n_layer, -1)\n",
    "\n",
    "    # Prefill\n",
    "    from suno_utils.gpt.bct.bct_generation_simple import prefill, BlockInputs, StreamState\n",
    "    prefill(model, sequence, kv_cache, block_table)\n",
    "    stream = StreamState(sequence, weight, kv_cache, block_table)\n",
    "\n",
    "    # Get the output block and prepare inputs\n",
    "    output_block = sequence[-1]\n",
    "    inputs = BlockInputs(output_block, cfg.block_size)\n",
    "\n",
    "    # Forward pass to get first logits\n",
    "    pos = stream.sequence.n_tokens - len(output_block) + len(inputs) - 1\n",
    "    last_token = inputs.last_token()\n",
    "    first_logits = model.forward(\n",
    "        last_token,\n",
    "        torch.tensor([pos]).cuda(),\n",
    "        cache_batch_idx=stream.block_table,\n",
    "        kv_cache=stream.kv_cache,\n",
    "    )\n",
    "\n",
    "    # Print the first text logits\n",
    "    print(\"First text logits shape:\", first_logits['text_output'].shape)\n",
    "    \n",
    "    # Convert to probabilities\n",
    "    text_probs = torch.nn.functional.softmax(first_logits['text_output'], dim=-1)\n",
    "    \n",
    "    # Get top-k predictions\n",
    "    top_k = 5\n",
    "    top_k_probs, top_k_indices = torch.topk(text_probs[0, 0], top_k)\n",
    "    print(f\"\\nTop {top_k} predictions after '{prefix}':\")\n",
    "    for i, (prob, idx) in enumerate(zip(top_k_probs, top_k_indices)):\n",
    "        token = model_container[\"tokenizer\"].decode([idx.item()])\n",
    "        print(f\"{i+1:2d}. Token: '{token}' (ID: {idx.item()}) - Probability: {prob.item():.4f}\")\n",
    "    \n",
    "    # Also show the raw logit values for comparison\n",
    "    top_k_logits = first_logits['text_output'][0, 0][top_k_indices]\n",
    "    print(f\"\\nCorresponding logit values:\")\n",
    "    for i, (logit, idx) in enumerate(zip(top_k_logits, top_k_indices)):\n",
    "        token = model_container[\"tokenizer\"].decode([idx.item()])\n",
    "        print(f\"{i+1:2d}. Token: '{token}' - Logit: {logit.item():.4f}\")\n",
    "    \n",
    "    return {\n",
    "        'prefix': prefix,\n",
    "        'logits': first_logits['text_output'][0, 0].cpu(),\n",
    "        'probs': text_probs[0, 0].cpu(),\n",
    "        'top_k_tokens': [model_container[\"tokenizer\"].decode([idx.item()]) for idx in top_k_indices],\n",
    "        'top_k_probs': top_k_probs.cpu(),\n",
    "        'top_k_logits': top_k_logits.cpu()\n",
    "    }\n",
    "\n",
    "# Test with different prefixes\n",
    "base_prompt = \"\"  # Your original prompt without any prefix\n",
    "prefixes_to_test = [\n",
    "    \"\",  # No prefix (original)\n",
    "    \"[no tags]\\n\",\n",
    "    \"[has tags]\\n\", \n",
    "]\n",
    "\n",
    "results = []\n",
    "for prefix in prefixes_to_test:\n",
    "    result = test_prompt_with_prefix(model, model_container, codes, cfg, base_prompt, prefix)\n",
    "    results.append(result)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "720044d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# meow"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}