Untitled.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "852385a9-bc7f-419e-8523-e3fa2c653166",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import tempfile\n",
    "import torchaudio\n",
    "import librosa\n",
    "from pydub import AudioSegment\n",
    "from flask import Flask, request, jsonify\n",
    "from flask_cors import CORS\n",
    "from pyngrok import ngrok\n",
    "from transformers import WhisperProcessor, WhisperForConditionalGeneration\n",
    "from werkzeug.utils import secure_filename\n",
    "from googletrans import Translator\n",
    "from huggingface_hub import login\n",
    "import logging\n",
    "import time\n",
    "\n",
    "# ---------------- CONFIG ----------------\n",
    "MODEL_NAME    = \"openai/whisper-medium\"\n",
    "MODEL_PATH    = \"best_model.pt\"\n",
    "UPLOAD_FOLDER = \"uploads\"\n",
    "HF_TOKEN      = \"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\"  # <---- Your token\n",
    "NGROK_TOKEN   = \"2w81ifj5xPmsRpAON5rvg4OlVxY_49NCKqCvAAgrJ2b8coFft\"  # <---- Your token\n",
    "\n",
    "os.makedirs(UPLOAD_FOLDER, exist_ok=True)\n",
    "\n",
    "# ----- LOGGING -----\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "# ----- AUTH -----\n",
    "login(HF_TOKEN)\n",
    "\n",
    "# ----- LOAD MODEL -----\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "logger.info(f\"Using device: {device}\")\n",
    "\n",
    "if device.type == \"cuda\":\n",
    "    torch.backends.cudnn.benchmark = True\n",
    "    os.environ[\"CUBLAS_WORKSPACE_CONFIG\"] = \":4096:8\"\n",
    "\n",
    "processor = WhisperProcessor.from_pretrained(\n",
    "    MODEL_NAME, use_auth_token=True, language=\"pashto\"\n",
    ")\n",
    "agri_tokens = [\"ګیدړه\", \"باغونه\", \"کانال\", \"زراعت\", \"مالداری\"]\n",
    "processor.tokenizer.add_tokens(agri_tokens)\n",
    "\n",
    "model = WhisperForConditionalGeneration.from_pretrained(\n",
    "    MODEL_NAME, use_auth_token=True\n",
    ")\n",
    "model.resize_token_embeddings(len(processor.tokenizer))\n",
    "model.load_state_dict(torch.load(MODEL_PATH, map_location=device))\n",
    "model.to(device).eval()\n",
    "\n",
    "translator = Translator()\n",
    "\n",
    "# ----- FLASK SETUP -----\n",
    "app = Flask(__name__)\n",
    "CORS(app)\n",
    "app.config[\"UPLOAD_FOLDER\"]      = UPLOAD_FOLDER\n",
    "app.config[\"MAX_CONTENT_LENGTH\"] = 10 * 1024 * 1024  # 10MB limit\n",
    "\n",
    "# ----- AUDIO PROCESSING -----\n",
    "def convert_to_wav(path_in: str) -> str:\n",
    "    \"\"\"\n",
    "    Convert any audio file to 16kHz mono WAV using:\n",
    "      1) TorchAudio\n",
    "      2) Pydub/FFmpeg\n",
    "      3) Librosa\n",
    "    \"\"\"\n",
    "    # 1) Try TorchAudio on the original file\n",
    "    try:\n",
    "        waveform, sample_rate = torchaudio.load(path_in)\n",
    "        logger.info(f\"Loaded with TorchAudio: {path_in}\")\n",
    "    except Exception as e:\n",
    "        logger.warning(f\"TorchAudio failed ({e}), trying Pydub/FFmpeg\")\n",
    "        # 2) Pydub fallback\n",
    "        try:\n",
    "            audio = AudioSegment.from_file(path_in)\n",
    "            audio = audio.set_frame_rate(16000).set_channels(1)\n",
    "            tmp = tempfile.NamedTemporaryFile(delete=False, suffix=\".wav\")\n",
    "            audio.export(\n",
    "                tmp.name,\n",
    "                format=\"wav\",\n",
    "                parameters=[\"-acodec\", \"pcm_s16le\"]\n",
    "            )\n",
    "            logger.info(f\"Pydub-converted audio saved to: {tmp.name}\")\n",
    "            return tmp.name\n",
    "        except Exception as e2:\n",
    "            logger.warning(f\"Pydub failed ({e2}), falling back to Librosa\")\n",
    "            # 3) Librosa fallback\n",
    "            y, sample_rate = librosa.load(path_in, sr=16000)\n",
    "            waveform = torch.from_numpy(y).unsqueeze(0)\n",
    "\n",
    "    else:\n",
    "        # Resample & mono mix if TorchAudio succeeded\n",
    "        if sample_rate != 16000:\n",
    "            resampler = torchaudio.transforms.Resample(\n",
    "                orig_freq=sample_rate,\n",
    "                new_freq=16000\n",
    "            )\n",
    "            waveform = resampler(waveform)\n",
    "        if waveform.shape[0] > 1:\n",
    "            waveform = torch.mean(waveform, dim=0, keepdim=True)\n",
    "\n",
    "    # Save final waveform as 16kHz PCM WAV\n",
    "    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=\".wav\")\n",
    "    torchaudio.save(\n",
    "        tmp.name,\n",
    "        waveform,\n",
    "        16000,\n",
    "        encoding=\"PCM_S\",\n",
    "        bits_per_sample=16\n",
    "    )\n",
    "    logger.info(f\"Final converted audio saved to: {tmp.name}\")\n",
    "    return tmp.name\n",
    "\n",
    "def transcribe_and_translate(wav_path: str):\n",
    "    \"\"\"\n",
    "    Always load the final WAV with librosa (avoids torchaudio backend issues),\n",
    "    then run through Whisper and Google Translate.\n",
    "    \"\"\"\n",
    "    # Load via librosa so we’re guaranteed to get a numpy array\n",
    "    y, sr = librosa.load(wav_path, sr=16000)\n",
    "    waveform = torch.from_numpy(y).unsqueeze(0).to(device)\n",
    "\n",
    "    inputs = processor(\n",
    "        audio=waveform.squeeze().cpu().numpy(),\n",
    "        sampling_rate=16000,\n",
    "        return_tensors=\"pt\"\n",
    "    ).to(device)\n",
    "\n",
    "    with torch.inference_mode():\n",
    "        ids = model.generate(inputs[\"input_features\"])\n",
    "        text = processor.batch_decode(ids, skip_special_tokens=True)[0]\n",
    "\n",
    "    try:\n",
    "        trans = translator.translate(text, dest=\"ur\")\n",
    "        urdu = trans.text\n",
    "    except Exception as e:\n",
    "        logger.error(f\"Translation error: {e}\")\n",
    "        urdu = \"Translation unavailable\"\n",
    "\n",
    "    return text, urdu\n",
    "\n",
    "# ----- ROUTES -----\n",
    "@app.route(\"/\", methods=[\"GET\"])\n",
    "def index():\n",
    "    return jsonify({\"status\": \"ok\"}), 200\n",
    "\n",
    "time.sleep(5)\n",
    "\n",
    "@app.route(\"/transcribe\", methods=[\"POST\"])\n",
    "def transcribe():\n",
    "    try:\n",
    "        if \"audio\" not in request.files:\n",
    "            return jsonify({\"error\": \"No audio file\"}), 400\n",
    "\n",
    "        f = request.files[\"audio\"]\n",
    "        if f.filename == \"\":\n",
    "            return jsonify({\"error\": \"Empty filename\"}), 400\n",
    "\n",
    "        filename = secure_filename(f.filename)\n",
    "        in_path = os.path.join(app.config[\"UPLOAD_FOLDER\"], filename)\n",
    "        f.save(in_path)\n",
    "\n",
    "        wav_path = None\n",
    "        try:\n",
    "            wav_path = convert_to_wav(in_path)\n",
    "            text, urdu = transcribe_and_translate(wav_path)\n",
    "        except Exception as e:\n",
    "            logger.error(f\"Audio processing error: {e}\")\n",
    "            return jsonify({\"error\": \"Audio processing failed\"}), 500\n",
    "        finally:\n",
    "            for p in (in_path, wav_path):\n",
    "                if p and os.path.exists(p):\n",
    "                    try:\n",
    "                        os.remove(p)\n",
    "                    except Exception as e:\n",
    "                        logger.warning(f\"Cleanup error: {e}\")\n",
    "\n",
    "        return jsonify({\n",
    "            \"pashto_transcription\": text,\n",
    "            \"urdu_translation\": urdu\n",
    "        }), 200\n",
    "\n",
    "    except Exception as e:\n",
    "        logger.error(f\"Endpoint error: {e}\")\n",
    "        return jsonify({\"error\": \"Server error\"}), 500\n",
    "\n",
    "# ----- RUN -----\n",
    "if __name__ == \"__main__\":\n",
    "    ngrok.set_auth_token(NGROK_TOKEN)\n",
    "    public_url = ngrok.connect(5000)\n",
    "    logger.info(f\"Ngrok tunnel -> {public_url}\")\n",
    "    app.run(host=\"0.0.0.0\", port=5000)\n"
   ]
  }
 ],
 "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}