{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "header"
      },
      "source": [
        "# 🚀 ConfereAI - Fast Training (GPU Edition)\n",
        "Este notebook permite treinar o motor neural do ConfereAI utilizando a GPU gratuita do Google Colab. \n",
        "\n",
        "**Instruções:**\n",
        "1. Vá em `Ambiente de Execução` > `Alterar tipo de ambiente` e selecione **T4 GPU**.\n",
        "2. Preencha as configurações abaixo.\n",
        "3. Execute as células em ordem."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "setup"
      },
      "outputs": [],
      "source": [
        "# @title 1. Instalar Dependências\n",
        "!pip install -q transformers[torch] librosa soundfile huggingface_hub accelerate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "config"
      },
      "outputs": [],
      "source": [
        "# @title 2. Configurações do Hugging Face\n",
        "HF_TOKEN = \"\" # @param {type:\"string\"}\n",
        "REPO_ID = \"TEDDyx86/confereai-dev\" # @param {type:\"string\"}\n",
        "BRANCH = \"main\" # @param {type:\"string\"}\n",
        "\n",
        "from huggingface_hub import HfApi, login\n",
        "if HF_TOKEN:\n",
        "    login(token=HF_TOKEN)\n",
        "else:\n",
        "    print(\"❌ Por favor, insira o seu HF_TOKEN!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "upload"
      },
      "outputs": [],
      "source": [
        "# @title 3. Upload do Dataset (.zip)\n",
        "from google.colab import files\n",
        "import zipfile\n",
        "import os\n",
        "import shutil\n",
        "\n",
        "uploaded = files.upload()\n",
        "dataset_zip = list(uploaded.keys())[0]\n",
        "\n",
        "DATASET_DIR = \"dataset_training\"\n",
        "if os.path.exists(DATASET_DIR): shutil.rmtree(DATASET_DIR)\n",
        "os.makedirs(DATASET_DIR)\n",
        "\n",
        "with zipfile.ZipFile(dataset_zip, 'r') as zip_ref:\n",
        "    zip_ref.extractall(DATASET_DIR)\n",
        "\n",
        "print(f\"✅ Dataset extraído em: {DATASET_DIR}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "training"
      },
      "outputs": [],
      "source": [
        "# @title 4. Executar Treinamento (Fine-Tuning)\n",
        "import torch\n",
        "from torch.utils.data import Dataset\n",
        "from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2ForSequenceClassification, Trainer, TrainingArguments\n",
        "import librosa\n",
        "\n",
        "BASE_MODEL = \"HyperMoon/wav2vec2-base-960h-finetuned-deepfake\"\n",
        "OUTPUT_DIR = \"local_finetuned_model\"\n",
        "\n",
        "class DeepfakeDataset(Dataset):\n",
        "    def __init__(self, root_dir, processor):\n",
        "        self.files = []\n",
        "        self.processor = processor\n",
        "        for label, folder in enumerate(['real', 'fake']):\n",
        "            path = os.path.join(root_dir, folder)\n",
        "            if os.path.exists(path):\n",
        "                for f in os.listdir(path):\n",
        "                    if f.endswith(('.wav', '.mp3', '.flac')):\n",
        "                        self.files.append({\"path\": os.path.join(path, f), \"label\": label})\n",
        "\n",
        "    def __len__(self): return len(self.files)\n",
        "    def __getitem__(self, idx):\n",
        "        item = self.files[idx]\n",
        "        speech, _ = librosa.load(item[\"path\"], sr=16000)\n",
        "        input_values = self.processor(speech, sampling_rate=16000, return_tensors=\"pt\", padding=\"max_length\", max_length=160000, truncation=True).input_values[0]\n",
        "        return {\"input_values\": input_values, \"labels\": torch.tensor(item[\"label\"], dtype=torch.long)}\n",
        "\n",
        "processor = Wav2Vec2FeatureExtractor.from_pretrained(BASE_MODEL)\n",
        "model = Wav2Vec2ForSequenceClassification.from_pretrained(BASE_MODEL, num_labels=2, ignore_mismatched_sizes=True)\n",
        "\n",
        "# Congelar base para focar no aprendizado das novas fraudes (Lógica Robusta)\n",
        "if hasattr(model, 'freeze_feature_extractor'):\n",
        "    model.freeze_feature_extractor()\n",
        "elif hasattr(model, 'freeze_feature_encoder'):\n",
        "    model.freeze_feature_encoder()\n",
        "\n",
        "if hasattr(model, 'wav2vec2'):\n",
        "    for param in model.wav2vec2.parameters(): param.requires_grad = False\n",
        "\n",
        "training_args = TrainingArguments(\n",
        "    output_dir=\"./results\",\n",
        "    num_train_epochs=5,\n",
        "    per_device_train_batch_size=4,\n",
        "    gradient_accumulation_steps=2,\n",
        "    learning_rate=2e-5,\n",
        "    logging_steps=1,\n",
        "    push_to_hub=False,\n",
        "    report_to=\"none\"\n",
        ")\n",
        "\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    train_dataset=DeepfakeDataset(DATASET_DIR, processor)\n",
        ")\n",
        "\n",
        "print(\"🚀 Iniciando treinamento na GPU...\")\n",
        "trainer.train()\n",
        "\n",
        "model.save_pretrained(OUTPUT_DIR)\n",
        "processor.save_pretrained(OUTPUT_DIR)\n",
        "print(f\"✅ Treinamento concluído. Modelo salvo em {OUTPUT_DIR}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "push"
      },
      "outputs": [],
      "source": [
        "# @title 5. Sincronizar com Hugging Face Space\n",
        "api = HfApi()\n",
        "print(f\"📦 Subindo modelo para {REPO_ID}...\")\n",
        "\n",
        "api.upload_folder(\n",
        "    folder_path=OUTPUT_DIR,\n",
        "    path_in_repo=OUTPUT_DIR,\n",
        "    repo_id=REPO_ID,\n",
        "    repo_type=\"space\",\n",
        "    token=HF_TOKEN,\n",
        "    commit_message=\"🤖 Auto-Update: Novo modelo treinado via Google Colab\"\n",
        ")\n",
        "\n",
        "print(\"✨ Sucesso! O seu Space irá reiniciar em breve com o novo modelo.\")"
      ]
    }
  ]
}