generation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c014cd14-5cb8-4515-b919-6852bf6c1e5e",
   "metadata": {},
   "source": [
    "# AudioGen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1a908c0-96c8-4a10-b164-94f83f086a9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from audiocraft.models import AudioGen\n",
    "\n",
    "model = AudioGen.get_pretrained('facebook/audiogen-medium')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f2666df-f8d0-46e5-99f9-ad70ba842668",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.set_generation_params(\n",
    "    use_sampling=True,\n",
    "    top_k=250,\n",
    "    duration=5\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b138165b-07f1-4ee6-8879-e4c46d79d7ef",
   "metadata": {},
   "source": [
    "## Audio Continuation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f07fec0e-bc9a-4c77-b88e-ef856f031b73",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torchaudio\n",
    "import torch\n",
    "from audiocraft.utils.notebook import display_audio\n",
    "\n",
    "def get_bip_bip(bip_duration=0.125, frequency=440,\n",
    "                duration=0.5, sample_rate=16000, device=\"cuda\"):\n",
    "    \"\"\"Generates a series of bip bip at the given frequency.\"\"\"\n",
    "    t = torch.arange(\n",
    "        int(duration * sample_rate), device=\"cuda\", dtype=torch.float) / sample_rate\n",
    "    wav = torch.cos(2 * math.pi * frequency * t)[None]\n",
    "    tp = (t % (2 * bip_duration)) / (2 * bip_duration)\n",
    "    envelope = (tp >= 0.5).float()\n",
    "    return wav * envelope"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cd26c1d-4368-4838-bd2a-392717247366",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we use a synthetic signal to prompt the generated audio.\n",
    "res = model.generate_continuation(\n",
    "    get_bip_bip(0.125).expand(2, -1, -1), \n",
    "    16000, ['Whistling with wind blowing', \n",
    "            'Typing on a typewriter'], \n",
    "    progress=True)\n",
    "display_audio(res, 16000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51f4041a-e164-455d-b366-e661ad880e3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can also use any audio from a file. Make sure to trim the file if it is too long!\n",
    "prompt_waveform, prompt_sr = torchaudio.load(\"Yun Hi Chala Chal.mp3\")\n",
    "# prompt_waveform, prompt_sr = torchaudio.load(\"../assets/sirens_and_a_humming_engine_approach_and_pass.mp3\")\n",
    "\n",
    "prompt_duration = 6\n",
    "prompt_waveform = prompt_waveform[..., :int(prompt_duration * prompt_sr)]\n",
    "output = model.generate_continuation(prompt_waveform, prompt_sample_rate=prompt_sr, progress=True)\n",
    "display_audio(output, sample_rate=16000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12fdf190-49bd-4ecd-9d13-afc54cf0a2b9",
   "metadata": {},
   "source": [
    "### Text-conditional Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5060c5c7-22bc-48b9-9c80-f93d74e322cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from audiocraft.utils.notebook import display_audio\n",
    "\n",
    "output = model.generate(\n",
    "    descriptions=[\n",
    "        'Ar Rahman music styled tabla',\n",
    "        # 'Subway train blowing its horn',\n",
    "        'A cat meowing',\n",
    "    ],\n",
    "    progress=True\n",
    ")\n",
    "display_audio(output, sample_rate=16000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2fb2439-a854-47ea-a149-bfbe76f09f50",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# MAGNeT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d32ad9a1-ab5e-448a-a485-fed5d4d5b3c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from audiocraft.models import MAGNeT\n",
    "\n",
    "model = MAGNeT.get_pretrained('facebook/magnet-small-10secs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf8a9599-5e87-49cc-b2ee-4cd0c8211da5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.set_generation_params(\n",
    "    use_sampling=True,\n",
    "    top_k=0,\n",
    "    top_p=0.9,\n",
    "    temperature=3.0,\n",
    "    max_cfg_coef=10.0,\n",
    "    min_cfg_coef=1.0,\n",
    "    decoding_steps=[int(20 * model.lm.cfg.dataset.segment_duration // 10),  10, 10, 10],\n",
    "    span_arrangement='stride1'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1211b5e9-14b9-4a01-83b2-20c07a0ac0c0",
   "metadata": {},
   "source": [
    "### Text-conditional Generation - Music"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39209539-eb4d-4cf2-a057-db032c196830",
   "metadata": {},
   "outputs": [],
   "source": [
    "from audiocraft.utils.notebook import display_audio\n",
    "\n",
    "###### Text-to-music prompts - examples ######\n",
    "text = \"80s bollywood music with deep instrumentals and base, hindi lyrics and very existential meaning\"\n",
    "# text = \"80s electronic track with melodic synthesizers, catchy beat and groovy bass. 170 bpm\"\n",
    "# text = \"Earthy tones, environmentally conscious, ukulele-infused, harmonic, breezy, easygoing, organic instrumentation, gentle grooves\"\n",
    "# text = \"Funky groove with electric piano playing blue chords rhythmically\"\n",
    "# text = \"Rock with saturated guitars, a heavy bass line and crazy drum break and fills.\"\n",
    "# text = \"A grand orchestral arrangement with thunderous percussion, epic brass fanfares, and soaring strings, creating a cinematic atmosphere fit for a heroic battle\"\n",
    "                   \n",
    "N_VARIATIONS = 3\n",
    "descriptions = [text for _ in range(N_VARIATIONS)]\n",
    "\n",
    "print(f\"text prompt: {text}\\n\")\n",
    "output = model.generate(descriptions=descriptions, progress=True, return_tokens=True)\n",
    "display_audio(output[0], sample_rate=model.compression_model.sample_rate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "769b0bc6-7d14-42a5-aa62-82a553ca9ed4",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Text-conditional Generation - Sound Effects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12efd8ff-fca8-4803-9f16-a1057e0766fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from audiocraft.models import MAGNeT\n",
    "\n",
    "model = MAGNeT.get_pretrained('facebook/audio-magnet-small')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f32b7c6d-c481-446b-bb0e-d918759276c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.set_generation_params(\n",
    "    use_sampling=True,\n",
    "    top_k=0,\n",
    "    top_p=0.8,\n",
    "    temperature=3.5,\n",
    "    max_cfg_coef=20.0,\n",
    "    min_cfg_coef=1.0,\n",
    "    decoding_steps=[int(20 * model.lm.cfg.dataset.segment_duration // 10),  10, 10, 10],\n",
    "    span_arrangement='stride1'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "396f7960-caf8-4155-abe2-82b8751f7c85",
   "metadata": {},
   "outputs": [],
   "source": [
    "from audiocraft.utils.notebook import display_audio\n",
    "               \n",
    "###### Text-to-audio prompts - examples ######\n",
    "text = \"Seagulls squawking as ocean waves crash while wind blows heavily into a microphone.\"\n",
    "# text = \"A toilet flushing as music is playing and a man is singing in the distance.\"\n",
    "\n",
    "N_VARIATIONS = 3\n",
    "descriptions = [text for _ in range(N_VARIATIONS)]\n",
    "\n",
    "print(f\"text prompt: {text}\\n\")\n",
    "output = model.generate(descriptions=descriptions, progress=True, return_tokens=True)\n",
    "display_audio(output[0], sample_rate=model.compression_model.sample_rate)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "945e88f1-b207-4e3e-b632-00fda63e401b",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# MusicGen\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ed9f9b2-994a-4172-8ec0-14c30faa3dda",
   "metadata": {},
   "outputs": [],
   "source": [
    "from audiocraft.models import MusicGen\n",
    "from audiocraft.models import MultiBandDiffusion\n",
    "\n",
    "USE_DIFFUSION_DECODER = True\n",
    "# Using small model, better results would be obtained with `medium` or `large`.\n",
    "model = MusicGen.get_pretrained('facebook/musicgen-small')\n",
    "if USE_DIFFUSION_DECODER:\n",
    "    mbd = MultiBandDiffusion.get_mbd_musicgen()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f4dd1ff-e4c4-4bac-923b-da668e9da2b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.set_generation_params(\n",
    "    use_sampling=True,\n",
    "    top_k=250,\n",
    "    duration=30\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a345572e-9378-461a-ab60-da6a5bf6b2aa",
   "metadata": {},
   "source": [
    "### Music Continuation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "527893fb-51db-4264-8d0b-7ae0a655398a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torchaudio\n",
    "import torch\n",
    "from audiocraft.utils.notebook import display_audio\n",
    "\n",
    "def get_bip_bip(bip_duration=0.125, frequency=440,\n",
    "                duration=0.5, sample_rate=32000, device=\"cuda\"):\n",
    "    \"\"\"Generates a series of bip bip at the given frequency.\"\"\"\n",
    "    t = torch.arange(\n",
    "        int(duration * sample_rate), device=\"cuda\", dtype=torch.float) / sample_rate\n",
    "    wav = torch.cos(2 * math.pi * 440 * t)[None]\n",
    "    tp = (t % (2 * bip_duration)) / (2 * bip_duration)\n",
    "    envelope = (tp >= 0.5).float()\n",
    "    return wav * envelope"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "815d241e-7cdb-4140-b87c-47d680d57def",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here we use a synthetic signal to prompt both the tonality and the BPM\n",
    "# of the generated audio.\n",
    "res = model.generate_continuation(\n",
    "    get_bip_bip(0.125).expand(2, -1, -1), \n",
    "    32000, ['Jazz jazz and only jazz', \n",
    "            'Hindi tabla with flute and sitar',\n",
    "            ], \n",
    "    progress=True)\n",
    "display_audio(res, 32000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1215becb-235f-4d96-858a-aa932360df2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can also use any audio from a file. Make sure to trim the file if it is too long!\n",
    "# prompt_waveform, prompt_sr = torchaudio.load(\"../assets/bach.mp3\")\n",
    "prompt_waveform, prompt_sr = torchaudio.load(\"Yun Hi Chala Chal.mp3\")\n",
    "prompt_duration = 8\n",
    "prompt_waveform = prompt_waveform[..., :int(prompt_duration * prompt_sr)]\n",
    "output = model.generate_continuation(prompt_waveform, prompt_sample_rate=prompt_sr, progress=True, return_tokens=True)\n",
    "display_audio(output[0], sample_rate=32000)\n",
    "if USE_DIFFUSION_DECODER:\n",
    "    out_diffusion = mbd.tokens_to_wav(output[1])\n",
    "    display_audio(out_diffusion, sample_rate=32000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b14b2c1-df02-46bb-925a-f3a0c2f402ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "from audiocraft.utils.notebook import display_audio\n",
    "\n",
    "output = model.generate(\n",
    "    descriptions=[\n",
    "        #'80s pop track with bassy drums and synth',\n",
    "        #'90s rock song with loud guitars and heavy drums',\n",
    "        #'Progressive rock drum and bass solo',\n",
    "        #'Punk Rock song with loud drum and power guitar',\n",
    "        #'Bluesy guitar instrumental with soulful licks and a driving rhythm section',\n",
    "        #'Jazz Funk song with slap bass and powerful saxophone',\n",
    "        # 'drum and bass beat with intense percussions',\n",
    "        'Hindi tabla with flute and sitar'\n",
    "    ],\n",
    "    progress=True, return_tokens=True\n",
    ")\n",
    "display_audio(output[0], sample_rate=32000)\n",
    "if USE_DIFFUSION_DECODER:\n",
    "    out_diffusion = mbd.tokens_to_wav(output[1])\n",
    "    display_audio(out_diffusion, sample_rate=32000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97b86e7d-86c8-4013-99c9-50662790d7b8",
   "metadata": {},
   "source": [
    "### Melody-conditional Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55ecd97c-1c52-4d82-b67e-03bffe9c6b18",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchaudio\n",
    "from audiocraft.utils.notebook import display_audio\n",
    "\n",
    "model = MusicGen.get_pretrained('facebook/musicgen-melody')\n",
    "model.set_generation_params(duration=8)\n",
    "\n",
    "melody_waveform, sr = torchaudio.load(\"Yun Hi Chala Chal.mp3\")\n",
    "melody_waveform = melody_waveform.unsqueeze(0).repeat(2, 1, 1)\n",
    "output = model.generate_with_chroma(\n",
    "    descriptions=[\n",
    "        # '80s pop track with bassy drums and synth',\n",
    "        '90s bollywood music with deep instrumentals and base, hindi lyrics and very existential meaning',\n",
    "        'Hindi tabla with flute and sitar'\n",
    "    ],\n",
    "    melody_wavs=melody_waveform,\n",
    "    melody_sample_rate=sr,\n",
    "    progress=True, return_tokens=True\n",
    ")\n",
    "display_audio(output[0], sample_rate=32000)\n",
    "if USE_DIFFUSION_DECODER:\n",
    "    out_diffusion = mbd.tokens_to_wav(output[1])\n",
    "    display_audio(out_diffusion, sample_rate=32000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a1a8f21-b5b6-418e-932b-ff17b19ff585",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Create Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afc1b231-cf97-4217-b392-6a77a9c0ba3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import csv\n",
    "import pandas as pd\n",
    "\n",
    "file_path = 'prompt_golden_data.csv'\n",
    "\n",
    "\n",
    "def csv_to_json(csv_file_path, json_file_path):\n",
    "    data = []\n",
    "\n",
    "    # Read the CSV file\n",
    "    with open(csv_file_path, 'r') as csv_file:\n",
    "        csv_reader = csv.DictReader(csv_file)\n",
    "\n",
    "        # Iterate over each row in the CSV file\n",
    "        for row in csv_reader:\n",
    "            data.append(row)\n",
    "\n",
    "    # Write the data to a JSON file\n",
    "    with open(json_file_path, 'w') as json_file:\n",
    "        json.dump(data, json_file, indent=4)\n",
    "\n",
    "    print(f\"CSV file '{csv_file_path}' has been converted to JSON file '{json_file_path}'.\")\n",
    "\n",
    "# Example usage\n",
    "csv_file_path = 'prompt_golden_data.csv'\n",
    "json_file_path = 'music_lyrics.json'\n",
    "csv_to_json(csv_file_path, json_file_path)\n",
    "# Step 1: Load the original JSON data\n",
    "with open('music_lyrics.json', 'r') as file:\n",
    "    data = json.load(file)\n",
    "\n",
    "# Step 2: Modify the data structure\n",
    "new_data = []\n",
    "\n",
    "for song in data:\n",
    "    new_data.append({\n",
    "        'instruction': 'Write a song lyric based on the given inputs and verse prompt.',\n",
    "        'input': \n",
    "            'genres:'+ song['genres']+\n",
    "            'progression:'+ song['progression']+\n",
    "            'start_key'+ song['start_key']+\n",
    "            'verse prompt'+ song['prompts']\n",
    "        ,\n",
    "        'output': \n",
    "            'lyrics'+ song['processed_lyrics']  # Ensure this key 'lyrics' or similar is what you want\n",
    "        \n",
    "    })\n",
    "\n",
    "# Step 3: Write the modified data to a new JSON file\n",
    "with open('restructured.json', 'w') as file:\n",
    "    json.dump(new_data, file, indent=4)\n",
    "\n",
    "\n",
    "\n",
    "with open('data/dataset_info.json', 'r') as file:\n",
    "    data = json.load(file)\n",
    "    data.append('music_lyrics_generation')\n",
    "    data['music_lyrics_generation']['file_name']='music_lyrics_generation.json'\n",
    "    data['music_lyrics_generation']['file_sha1']='7df69e4325ad88feef052b3c086b4434867b120a'\n",
    "    json.dump(data, file, indent=4)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd6b5876-bffb-42d3-8e64-32d4bb605396",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Create melody of generated lyrics "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe0ecbd0-895f-4332-8d1c-b4a30c9e04a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "sys.path.append(\"../\")\n",
    "import yaml\n",
    "import json\n",
    "import numpy as np\n",
    "import torch\n",
    "import math\n",
    "import copy\n",
    "import re\n",
    "from transformers import AutoModel, AutoTokenizer\n",
    "# ckpt_path = '/Mar2Ding/songcomposer_pretrain'\n",
    "ckpt_path = 'Mar2Ding/songcomposer_sft'\n",
    "tokenizer = AutoTokenizer.from_pretrained(ckpt_path, trust_remote_code=True)\n",
    "model = AutoModel.from_pretrained(ckpt_path, trust_remote_code=True).cuda().half()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95fab8d8-3c46-48da-9252-935853059273",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = 'Compose a tune in harmony with the accompanying lyrics. <bol> Total 6 lines.\\\n",
    "The first line: NLP, the class that sets us free\\n\\\n",
    "The second line: Prof. Srihari and Sayantal, our guiding team\\n\\\n",
    "The third line: Natural Language Processing, oh so fine They want to grade us high, all the time\\n\\\n",
    "The fourth line: Parsing, disambiguating, we’re on a roll NLP, you’re in our soul\\n\\\n",
    "The fifth line: Natural Language Processing, oh so fine They want to grade us high, all the time\\n\\\n",
    "The sixth line: Prof. Srihari, Sayantal Pal, we thank you NLP, our passion true\\n<eol>'\n",
    "####### m2l #######\n",
    "model.inference(prompt, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c31147ff-f3f7-42a6-aedb-3ddc11ead0b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "line = 'The first line:<E4> , <154> , <88> |<E4> , <134> , <88> |<E4> , <137> , <79> |<F#4> , <151> , <79> |<E4> , <154> , <79> |<D#4> , <154> , <79> | <C#4> , <157> , <79> | <B3> , <172> , <127> The second line:<E4> , <151> , <88> |<E4> , <137> , <88> |<E4> , <137> , <79> |<F#4> , <151> , <79> |<E4> , <151> , <79> |<D#4> , <160> , <79> |<C#4> , <157> , <79> The third line:<B3> , <151> , <79> |<G#3> , <137> , <79> |<B3> , <151> , <79> |<G#3> , <189> , <79> |<F#3> , <157> , <79> |<G#3> , <137> , <79> The fourth line:<G#3> , <147> , <79> |<F#3> , <144> , <79> |<E3> , <151> , <79> |<F#3> , <141> , <79> |<G#3> , <166> , <79> |<B3> , <219> , <160> The fifth line:<E4> , <154> , <88> |<E4> , <130> , <88> |<E4> , <144> , <79> |<F#4> , <147> , <79> |<E4> , <157> , <79> |<D#4> , <154> , <79> |<C#4> , <151> , <79> |<B3> , <118> , <79> |<B3> , <118> , <79> |<G#3> , <207> , <79> |<B3> , <205> , <79> |的, <G#3> , <205> , <79>'\n",
    "from finetune.utils import gen_midi\n",
    "gen_midi(line, 'text')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "853b2019-1491-492f-b6df-1a17ef2dc0c9",
   "metadata": {},
   "source": [
    "# Finetune LLama 3 model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9005f0be-6426-4c7c-893e-34ee8d40a328",
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /content/\n",
    "%rm -rf LLaMA-Factory\n",
    "!git clone https://github.com/hiyouga/LLaMA-Factory.git\n",
    "%cd LLaMA-Factory\n",
    "%ls\n",
    "!pip install \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\"\n",
    "!pip install --no-deps xformers==0.0.25\n",
    "!pip install .[bitsandbytes]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d633a6a-7240-47ea-866a-7ff74b02630b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "try:\n",
    "  assert torch.cuda.is_available() is True\n",
    "except AssertionError:\n",
    "  print(\"Please set up a GPU before using LLaMA Factory\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f42ba82-7e5a-40ae-afe7-fa759ce13c56",
   "metadata": {},
   "source": [
    "## Fine-tune model via Command Line\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb79abbb-73bc-4616-b10d-f6ce94d3921c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd /content/LLaMA-Factory/\n",
    "!GRADIO_SHARE=1 llamafactory-cli webui\n",
    "\n",
    "\n",
    "import json\n",
    "\n",
    "args = dict(\n",
    "  stage=\"sft\",                        # do supervised fine-tuning\n",
    "  do_train=True,\n",
    "  model_name_or_path=\"unsloth/llama-3-8b-Instruct-bnb-4bit\", # use bnb-4bit-quantized Llama-3-8B-Instruct model\n",
    "  dataset=\"identity,alpaca_gpt4_en,\",             # use alpaca and identity datasets\n",
    "  template=\"llama3\",                     # use llama3 prompt template\n",
    "  finetuning_type=\"lora\",                   # use LoRA adapters to save memory\n",
    "  lora_target=\"all\",                     # attach LoRA adapters to all linear layers\n",
    "  output_dir=\"llama3_lora\",                  # the path to save LoRA adapters\n",
    "  per_device_train_batch_size=2,               # the batch size\n",
    "  gradient_accumulation_steps=4,               # the gradient accumulation steps\n",
    "  lr_scheduler_type=\"cosine\",                 # use cosine learning rate scheduler\n",
    "  logging_steps=10,                      # log every 10 steps\n",
    "  warmup_ratio=0.1,                      # use warmup scheduler\n",
    "  save_steps=1000,                      # save checkpoint every 1000 steps\n",
    "  learning_rate=5e-5,                     # the learning rate\n",
    "  num_train_epochs=3.0,                    # the epochs of training\n",
    "  max_samples=500,                      # use 500 examples in each dataset\n",
    "  max_grad_norm=1.0,                     # clip gradient norm to 1.0\n",
    "  quantization_bit=4,                     # use 4-bit QLoRA\n",
    "  loraplus_lr_ratio=16.0,                   # use LoRA+ algorithm with lambda=16.0\n",
    "  use_unsloth=True,                      # use UnslothAI's LoRA optimization for 2x faster training\n",
    "  fp16=True,                         # use float16 mixed precision training\n",
    ")\n",
    "\n",
    "json.dump(args, open(\"train_llama3.json\", \"w\", encoding=\"utf-8\"), indent=2)\n",
    "\n",
    "%cd /content/LLaMA-Factory/\n",
    "\n",
    "!llamafactory-cli train train_llama3.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "463f4c2c-be33-4690-9c70-39e9f8606799",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json \n",
    "\n",
    "with open(\"data/music_lyrics_generation.json\", \"r\", encoding=\"utf-8\") as f:\n",
    "  dataset = json.load(f)\n",
    "  print(dataset[:1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd123326-b697-4c01-a2d4-b28d212eadb8",
   "metadata": {},
   "source": [
    "### Infer the fine-tuned model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ebafb40-ce24-4e70-b879-d2b63516d0c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llmtuner.chat import ChatModel\n",
    "from llmtuner.extras.misc import torch_gc\n",
    "\n",
    "%cd /content/LLaMA-Factory/\n",
    "\n",
    "args = dict(\n",
    "  model_name_or_path=\"unsloth/llama-3-8b-Instruct-bnb-4bit\", # use bnb-4bit-quantized Llama-3-8B-Instruct model\n",
    "  adapter_name_or_path=\"llama3_lora\",            # load the saved LoRA adapters\n",
    "  template=\"llama3\",                     # same to the one in training\n",
    "  finetuning_type=\"lora\",                  # same to the one in training\n",
    "  quantization_bit=4,                    # load 4-bit quantized model\n",
    "  use_unsloth=True,                     # use UnslothAI's LoRA optimization for 2x faster generation\n",
    ")\n",
    "chat_model = ChatModel(args)\n",
    "\n",
    "messages = []\n",
    "print(\"Welcome to the CLI application, use `clear` to remove the history, use `exit` to exit the application.\")\n",
    "while True:\n",
    "  query = input(\"\\nUser: \")\n",
    "  if query.strip() == \"exit\":\n",
    "    break\n",
    "  if query.strip() == \"clear\":\n",
    "    messages = []\n",
    "    torch_gc()\n",
    "    print(\"History has been removed.\")\n",
    "    continue\n",
    "\n",
    "  messages.append({\"role\": \"user\", \"content\": query})\n",
    "  print(\"Assistant: \", end=\"\", flush=True)\n",
    "\n",
    "  response = \"\"\n",
    "  for new_text in chat_model.stream_chat(messages):\n",
    "    print(new_text, end=\"\", flush=True)\n",
    "    response += new_text\n",
    "  print()\n",
    "  messages.append({\"role\": \"assistant\", \"content\": response})\n",
    "\n",
    "torch_gc()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c478dbf7-5183-4aae-aeab-030471b4f51d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llmtuner import ChatModel\n",
    "from llmtuner.extras.misc import torch_gc\n",
    "\n",
    "\n",
    "chat_model = ChatModel(dict(\n",
    "  model_name_or_path=\"unsloth/llama-3-8b-Instruct-bnb-4bit\", # use bnb-4bit-quantized Llama-3-8B-Instruct model\n",
    "  adapter_name_or_path=\"llama3_lora\",            # load the saved LoRA adapters\n",
    "  finetuning_type=\"lora\",                  # same to the one in training\n",
    "  template=\"llama3\",                     # same to the one in training\n",
    "  quantization_bit=4,                    # load 4-bit quantized model\n",
    "  use_unsloth=True,                     # use UnslothAI's LoRA optimization for 2x faster generation\n",
    "))\n",
    "\n",
    "messages = []\n",
    "while True:\n",
    "  query = input(\"\\nUser: \")\n",
    "  if query.strip() == \"exit\":\n",
    "    break\n",
    "\n",
    "  if query.strip() == \"clear\":\n",
    "    messages = []\n",
    "    torch_gc()\n",
    "    print(\"History has been removed.\")\n",
    "    continue\n",
    "\n",
    "    \n",
    "\n",
    "    query = \"'instruction': 'Write a song lyric based on the given inputs and verse prompt.'\n",
    "\n",
    "    'user': 'genres': ['canadian pop', 'pop', 'post-teen pop']\n",
    "            'progression':['A', 'Em', 'G']\n",
    "            'start_key': 'Bm'\n",
    "            'verse prompt': 'I was fifteen when the world put me on a pedestal and told me Im the best.'\n",
    "    \"\n",
    "    messages.append({\"role\": \"user\", \"content\": query})     # add query to messages\n",
    "    print(\"Assistant: \", end=\"\", flush=True)\n",
    "    response = \"\"\n",
    "    for new_text in chat_model.stream_chat(messages):      # stream generation\n",
    "      print(new_text, end=\"\", flush=True)\n",
    "      response += new_text\n",
    "    print()\n",
    "    messages.append({\"role\": \"assistant\", \"content\": response}) # add response to messages\n",
    "\n",
    "torch_gc()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}