{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\"\n",
    "import sys\n",
    "import numpy as np\n",
    "sys.path.append(\"/home/christian/code/christian/scripts\")  # Point to the scripts directory\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from typing import List, Optional, Tuple, Dict, Any\n",
    "from train_gpt import GPTModel, GPTConfig, load_model, load_tokenizer, prepare_text_inference\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "from typing import List, Optional, Tuple, Dict, Any\n",
    "\n",
    "def generate_semantic_tokens(\n",
    "    model: torch.nn.Module,\n",
    "    text: str,\n",
    "    tokenizer,\n",
    "    config: Any,\n",
    "    max_length: int = 750,\n",
    "    temperature: float = 1.0,\n",
    "    top_k: int = 0,\n",
    "    top_p: float = 0.9,\n",
    "    repetition_penalty: float = 1.0,\n",
    "    do_sample: bool = True,\n",
    "    device: str = \"cuda\",\n",
    ") -> np.ndarray:\n",
    "    \"\"\"\n",
    "    Generate semantic tokens from text input using the trained GPT model.\n",
    "    \n",
    "    Args:\n",
    "        model: The trained GPT model\n",
    "        text: Input text (can include tags in format \"[tag1, tag2]\" followed by lyrics)\n",
    "        tokenizer: Tokenizer for encoding text\n",
    "        config: Model configuration\n",
    "        max_length: Maximum number of semantic tokens to generate\n",
    "        temperature: Sampling temperature (higher = more random)\n",
    "        top_k: Number of highest probability tokens to keep (0 = disable)\n",
    "        top_p: Cumulative probability threshold for nucleus sampling\n",
    "        repetition_penalty: Penalty for repeating tokens\n",
    "        do_sample: Whether to sample or use greedy decoding\n",
    "        device: Device to use for inference (\"cuda\" or \"cpu\")\n",
    "        \n",
    "    Returns:\n",
    "        Generated semantic tokens as numpy array\n",
    "    \"\"\"\n",
    "    model.eval()\n",
    "    \n",
    "    # Process input text\n",
    "    if isinstance(text, str):\n",
    "        # Extract tags and lyrics if they exist\n",
    "        tags = []\n",
    "        lyrics = text\n",
    "        if text.startswith(\"[\") and \"]\" in text:\n",
    "            tag_section, rest = text.split(\"]\", 1)\n",
    "            tags = [tag.strip() for tag in tag_section[1:].split(\",\")]\n",
    "            lyrics = rest.strip()\n",
    "            \n",
    "        # Prepare text with proper formatting\n",
    "        text = prepare_text_inference(tags, lyrics)\n",
    "    \n",
    "    # Tokenize input text\n",
    "    text_codes = tokenizer.encode(text).ids\n",
    "    # Add the infer token at the end\n",
    "    text_input_ids = torch.tensor(text_codes).unsqueeze(0).to(device)\n",
    "    text_len = text_input_ids.size(1)\n",
    "    \n",
    "    # Prepare for tracking generated tokens\n",
    "    generated_tokens = []\n",
    "    \n",
    "    # Initialize an empty list for semantic tokens\n",
    "    semantic_tokens_list = []\n",
    "    \n",
    "    # Convert model to appropriate precision\n",
    "    dtype = list(model.parameters())[0].dtype\n",
    "    model = model.to(device=device, dtype=dtype)\n",
    "    \n",
    "    # Generate tokens one by one\n",
    "    with torch.no_grad():\n",
    "        for i in range(max_length):\n",
    "            # For each step, create semantic input tensor from generated tokens\n",
    "            if len(semantic_tokens_list) > 0:\n",
    "                semantic_input_ids = torch.tensor([semantic_tokens_list], device=device)\n",
    "            else:\n",
    "                # For the first step, start with the semantic SOS token\n",
    "                semantic_input_ids = torch.tensor([[config.semantic_sos_token]], dtype=torch.long, device=device)\n",
    "                semantic_tokens_list.append(config.semantic_sos_token)\n",
    "            # Create input by concatenating text and semantic tokens\n",
    "            combined_input_ids = torch.cat([text_input_ids, semantic_input_ids], dim=1)\n",
    "            \n",
    "            # Create attention mask for the full sequence\n",
    "            # This is critical to ensure dimensions match in the model\n",
    "            attention_mask = torch.ones(combined_input_ids.size(), device=device).bool()\n",
    "            \n",
    "            # Forward pass for token generation\n",
    "            with torch.autocast(device_type=\"cuda\" if device == \"cuda\" else \"cpu\", dtype=torch.bfloat16 if dtype == torch.bfloat16 else torch.float32):\n",
    "                # The forward pass expects separate text and semantic inputs\n",
    "                logits = model(\n",
    "                    text_input_ids=text_input_ids,\n",
    "                    semantic_input_ids=semantic_input_ids,\n",
    "                    attention_mask=attention_mask\n",
    "                )\n",
    "            \n",
    "            # Get the last token's logits for prediction (the next token)\n",
    "            # If we have semantic tokens, look at the last semantic token\n",
    "            # Otherwise, look at the last text token\n",
    "            next_token_logits = logits[:, -1, :]\n",
    "            \n",
    "            # Apply temperature\n",
    "            next_token_logits = next_token_logits / max(temperature, 1e-8)\n",
    "            \n",
    "            # Apply repetition penalty\n",
    "            if repetition_penalty != 1.0 and len(semantic_tokens_list) > 0:\n",
    "                for prev_token in set(semantic_tokens_list):\n",
    "                    next_token_logits[:, prev_token] /= repetition_penalty\n",
    "            \n",
    "            # Apply top-k filtering\n",
    "            if top_k > 0:\n",
    "                indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]\n",
    "                next_token_logits[indices_to_remove] = -float('Inf')\n",
    "            \n",
    "            # Apply top-p (nucleus) filtering\n",
    "            if 0.0 < top_p < 1.0:\n",
    "                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)\n",
    "                cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)\n",
    "                \n",
    "                # Remove tokens with cumulative probability above the threshold\n",
    "                sorted_indices_to_remove = cumulative_probs > top_p\n",
    "                # Shift the indices to the right to keep the first token above threshold\n",
    "                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()\n",
    "                sorted_indices_to_remove[..., 0] = 0\n",
    "                \n",
    "                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)\n",
    "                next_token_logits[indices_to_remove] = -float('Inf')\n",
    "            \n",
    "            # Sample from the filtered distribution\n",
    "            if do_sample:\n",
    "                probs = F.softmax(next_token_logits, dim=-1)\n",
    "                next_token = torch.multinomial(probs, num_samples=1).item()\n",
    "            else:\n",
    "                # Greedy decoding\n",
    "                next_token = torch.argmax(next_token_logits, dim=-1).item()\n",
    "            \n",
    "            # Check if we've reached the EOS token\n",
    "            if next_token == config.semantic_eos_token:\n",
    "                break\n",
    "            \n",
    "            # Add the token to our generated sequence\n",
    "            semantic_tokens_list.append(next_token)\n",
    "    \n",
    "    # Convert to numpy array\n",
    "    return np.array(semantic_tokens_list)\n",
    "\n",
    "def load_model_checkpoint(checkpoint_path, device=\"cuda\"):\n",
    "    \"\"\"\n",
    "    Load a trained GPT model from a checkpoint.\n",
    "    \n",
    "    Args:\n",
    "        checkpoint_path: Path to the checkpoint file\n",
    "        device: Device to load the model on\n",
    "        \n",
    "    Returns:\n",
    "        model: The loaded model\n",
    "        config: The model configuration\n",
    "    \"\"\"\n",
    "    checkpoint = torch.load(checkpoint_path, map_location=device)\n",
    "    \n",
    "    # Get config from checkpoint or create default\n",
    "    if \"config\" in checkpoint:\n",
    "        config = checkpoint[\"config\"]\n",
    "    else:\n",
    "        config = GPTConfig()\n",
    "    \n",
    "    # Initialize model with config\n",
    "    model = GPTModel(config)\n",
    "    \n",
    "    # Load state dict\n",
    "    if \"model\" in checkpoint:\n",
    "        model.load_state_dict(checkpoint[\"model\"])\n",
    "    else:\n",
    "        model.load_state_dict(checkpoint)\n",
    "    \n",
    "    model = model.to(device)\n",
    "    model.eval()\n",
    "    \n",
    "    return model, config\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = load_tokenizer()\n",
    "model, config = load_model(\"/app/suno/christian/checkpoints/gpt/2025-04-01_18-30-50_s6899/last_ckpt.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Example text with tags and lyrics\n",
    "text = \"[pop, upbeat, happy] This is a sample lyric for a happy pop song\"\n",
    "\n",
    "# Generate semantic tokens\n",
    "semantic_tokens = generate_semantic_tokens(\n",
    "    model=model,\n",
    "    text=text,\n",
    "    tokenizer=tokenizer,\n",
    "    config=config,\n",
    "    max_length=750,\n",
    "    temperature=0.8,\n",
    "    #top_p=0.9,\n",
    "    #repetition_penalty=1.1,\n",
    ")\n",
    "\n",
    "print(f\"Generated {len(semantic_tokens)} semantic tokens\")\n",
    "print(f\"Sample: {semantic_tokens}\")  # Print first 20 tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(len(semantic_tokens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save tokens to npz\n",
    "np.savez(\"/home/christian/code/christian/notebooks/outputs/semantic_codes/generated_semantic.npz\", semantic_codes=semantic_tokens)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "suno_diff",
   "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.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}