voice_crux_backend.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vPgu04Smnq3M"
      },
      "outputs": [],
      "source": [
        "%%writefile diarization.py\n",
        "\n",
        "import whisperx\n",
        "import gc\n",
        "import pandas as pd\n",
        "import base64\n",
        "\n",
        "def diarize(base64_audio, number):\n",
        "  device = \"cuda\"\n",
        "  batch_size = 4\n",
        "  compute_type = \"float16\"\n",
        "\n",
        "  audio_bytes = base64.b64decode(base64_audio)\n",
        "  with open(\"output.wav\", \"wb\") as wav_file:\n",
        "    wav_file.write(audio_bytes)\n",
        "\n",
        "  model = whisperx.load_model(\"large-v2\", device, compute_type=compute_type)\n",
        "  audio = whisperx.load_audio(\"output.wav\")\n",
        "  result1 = model.transcribe(audio, batch_size=batch_size)\n",
        "  model_a, metadata = whisperx.load_align_model(language_code=result1[\"language\"], device=device)\n",
        "  result = whisperx.align(result1[\"segments\"], model_a, metadata, audio, device, return_char_alignments=False)\n",
        "  diarize_model = whisperx.DiarizationPipeline(use_auth_token=\"HF_TOKEN\", device=device)\n",
        "  diarize_segments = diarize_model(audio, min_speakers=number, max_speakers=number)\n",
        "  result2 = whisperx.assign_word_speakers(diarize_segments, result)\n",
        "\n",
        "  return result2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "723UE_yKdIdG"
      },
      "outputs": [],
      "source": [
        "%%writefile summarization.py\n",
        "\n",
        "import google.generativeai as genai\n",
        "\n",
        "genai.configure(api_key=\"API_KEY\")\n",
        "\n",
        "SYSTEM_PROMPT = (\n",
        "    \"YOU ARE A SUMMARIZATION ASSISTANT. YOUR TASK IS TO:\"\n",
        "\n",
        "        \"1. READ THE PROVIDED TEXT CAREFULLY.\"\n",
        "        \"2. SUMMARIZE IT INTO A MAXIMUM OF 200 WORDS.\"\n",
        "        \"3. IF YOU THINK THAT THE INPUT TEXT IS TOO SMALL FOR SUMMARIZATION DONT OUTPUT UNNECESSARY TEXT OR WARNINGS OR ERRORS RATHER JUST OUTPUT THE ORIGINAL TEXT, WITHOUT TAMPERING\"\n",
        "\n",
        "    \"WRITE THE SUMMARY IN AN ABSTRACT STYLE, ENSURING THE CORE MESSAGE AND KEY IDEAS ARE RETAINED BUT WITHOUT COPYING ANY LINES VERBATIM FROM THE INPUT TEXT.\"\n",
        "    \"OUTPUT RULES:\"\n",
        "\n",
        "        \"1. RETURN ONLY THE SUMMARY.\"\n",
        "        \"2. DO NOT INCLUDE ANY INTRODUCTIONS, EXPLANATIONS, OR FORMATTING IN YOUR RESPONSE.\"\n",
        "        \"3. THE OUTPUT MUST CONSIST ONLY OF THE ABSTRACT SUMMARY.\"\n",
        ")\n",
        "\n",
        "def summarize(text):\n",
        "    prompt = f\"\"\"### System Prompt:\n",
        "    {SYSTEM_PROMPT}\n",
        "\n",
        "    ### User Prompt:\n",
        "    Text to be SUMMARIZED: {text}\n",
        "\n",
        "    ### Assistant Response:\"\"\"\n",
        "\n",
        "    model = genai.GenerativeModel(\"gemini-1.5-flash\")\n",
        "    response = model.generate_content(prompt)\n",
        "    return response.text.replace(\"  \",\" \").strip().strip('\\n')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "EXVGwsNRhp0z"
      },
      "outputs": [],
      "source": [
        "!pip install whisperx pyngrok flask"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zb7wegUs8huE"
      },
      "outputs": [],
      "source": [
        "from pyngrok import ngrok\n",
        "from flask import Flask, request\n",
        "from diarization import diarize\n",
        "from summarization import summarize"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3eyrvuiBvPma"
      },
      "outputs": [],
      "source": [
        "port_no = 5000\n",
        "app = Flask(__name__)\n",
        "ngrok.set_auth_token(\"NGROK_AUTH_TOKEN\")\n",
        "public_url = \"NGROK_PUBLIC_URL\"\n",
        "\n",
        "@app.route(\"/generate\", methods=[\"POST\"])\n",
        "def generate():\n",
        "  base64_audio = request.json[\"audio\"]\n",
        "  number = request.json[\"number\"]\n",
        "  diarized_data = diarize(base64_audio,number)\n",
        "\n",
        "  sum_dict={}\n",
        "\n",
        "  for i in range(number):\n",
        "    con_text=\"\"\n",
        "    for j in range (len(diarized_data[\"segments\"])):\n",
        "      if i <= 9:\n",
        "        num = '0' + str(i)\n",
        "      else:\n",
        "        num = str(i)\n",
        "\n",
        "      if diarized_data[\"segments\"][j][\"speaker\"] == \"SPEAKER_\"+num:\n",
        "        con_text = con_text + \" \" + diarized_data[\"segments\"][j][\"text\"]\n",
        "\n",
        "    sum_dict[\"SPEAKER_\"+num] = summarize(con_text)\n",
        "\n",
        "  return {\"d_data\":diarized_data, \"summary\":sum_dict}, 201\n",
        "\n",
        "ngrok_tunnel = ngrok.connect(port_no, bind_tls=True, hostname=\"NGROK_HOSTNAME\")\n",
        "if __name__ == '__main__':\n",
        "  print(f\"Ngrok tunnel: {ngrok_tunnel.public_url}\")\n",
        "  app.run(host=\"0.0.0.0\", debug=True, port=port_no, use_reloader=False)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}