Upload hexa_colab.ipynb with huggingface_hub

hexa_colab.ipynb

@@ -0,0 +1,165 @@
+{
+    "cells": [
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "# Hexa TTS - Free Colab Training (15GB GPU Optimized)\n",
+                "\n",
+                "**Compatibility:** Verified for **Tesla T4 (15GB usable VRAM)**.\\n",
+                "**Model Config:** Hexa-Base (~350M params).\\n",
+                "\n",
+                "## Setup\n",
+                "1.  Set Runtime to **T4 GPU**.\n",
+                "2.  Run all cells."
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# 1. Install Dependencies\n",
+                "!pip install -q -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118\n",
+                "!pip install -q -U transformers accelerate peft bitsandbytes soundfile phonemizer einops"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# 2. Clone Your Repository\n",
+                "import os\n",
+                "!git clone https://huggingface.co/Hexa09/hexa-tts-5b\n",
+                "root_dir = \"/content/hexa-tts-5b\"\n",
+                "os.chdir(root_dir)\n",
+                "\n",
+                "# Fix paths for Colab\n",
+                "import sys\n",
+                "sys.path.append(root_dir)"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# 3. System Imports\n",
+                "import torch\n",
+                "import gc\n",
+                "from transformers import Trainer, TrainingArguments\n",
+                "from src.hf_model import HexaModel, HexaHFConfig\n",
+                "from src.config import HexaConfig\n",
+                "from src.dataset_clean import HexaDataset, collate_fn\n",
+                "from get_data import download_data\n",
+                "\n",
+                "# Clear RAM\n",
+                "gc.collect()\n",
+                "torch.cuda.empty_cache()"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# 4. Generate Synthetic Data\n",
+                "if not os.path.exists(\"./data/metadata.csv\"):\n",
+                "    print(\"Generating synthetic data...\")\n",
+                "    download_data()"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# 5. Initialize Model (15GB Safe Config)\n",
+                "print(\"Initializing Hexa-Base (350M)...\")\n",
+                "\n",
+                "# 350M Params = ~700MB VRAM (Weights)\n",
+                "# Full Training State = ~5GB VRAM\n",
+                "# This leaves ~10GB headroom on a 15GB card.\n",
+                "\n",
+                "hexa_conf = HexaConfig(\n",
+                "    dim=1024,      # Optimized\n",
+                "    depth=24,      # Optimized\n",
+                "    heads=16,      # Optimized\n",
+                "    dim_head=64    \n",
+                ")\n",
+                "hf_config = HexaHFConfig(**hexa_conf.__dict__)\n",
+                "\n",
+                "# Initialize directly on GPU\n",
+                "with torch.device(\"cuda\"):\n",
+                "    model = HexaModel(hf_config)\n",
+                "\n",
+                "# Move to FP16 \n",
+                "model = model.half()\n",
+                "model.gradient_checkpointing_enable()\n",
+                "model.enable_input_require_grads()\n",
+                "\n",
+                "print(f\"Model Ready. Parameters: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M\")"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# 6. Training Arguments (Safe Mode)\n",
+                "args = TrainingArguments(\n",
+                "    output_dir=\"./hexa_colab_checkpoints\",\n",
+                "    per_device_train_batch_size=2,  # Reduced to 2 for 15GB safety\n",
+                "    gradient_accumulation_steps=8,  # Increased to maintain effective batch size\n",
+                "    learning_rate=2e-4,\n",
+                "    num_train_epochs=3,\n",
+                "    logging_steps=1,\n",
+                "    fp16=True,                      \n",
+                "    gradient_checkpointing=True,    \n",
+                "    report_to=\"none\",\n",
+                "    dataloader_num_workers=0\n",
+                ")\n",
+                "\n",
+                "dataset = HexaDataset(\"./data\", hexa_conf)\n",
+                "\n",
+                "trainer = Trainer(\n",
+                "    model=model,\n",
+                "    args=args,\n",
+                "    train_dataset=dataset,\n",
+                "    data_collator=collate_fn,\n",
+                ")\n",
+                "\n",
+                "print(\"Starting Training...\")\n",
+                "trainer.train()"
+            ]
+        }
+    ],
+    "metadata": {
+        "kernelspec": {
+            "display_name": "Python 3",
+            "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.10.12"
+        }
+    },
+    "nbformat": 4,
+    "nbformat_minor": 5
+}