transcription.py

import subprocess
import torch
from transformers import pipeline
from logging_config import logger, log_buffer

device = "cuda" if torch.cuda.is_available() else "cpu"


def convert_audio_to_wav(input_file: str, output_file: str, ffmpeg_path: str) -> str:
    logger.info(f"Converting {input_file} to WAV format: {output_file}")
    cmd = [
        ffmpeg_path,
        "-y",  # Overwrite output files without asking
        "-i", input_file,
        "-ar", "16000",  # Set audio sampling rate to 16kHz
        "-ac", "1",      # Set number of audio channels to mono
        output_file
    ]
    try:
        subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        logger.info("Audio conversion to WAV completed successfully.")
        return output_file
    except subprocess.CalledProcessError as e:
        ffmpeg_error = e.stderr.decode()
        logger.error(f"ffmpeg error: {ffmpeg_error}")
        raise RuntimeError("Failed to convert audio to WAV.") from e


def run_whisper_transcription(wav_file_path: str, device: str):
    try:
        asr_pipeline = pipeline(
            "automatic-speech-recognition",
            model="openai/whisper-small",
            device=0 if device == "cuda" else -1,
            return_timestamps=True,
            generate_kwargs={"task": "transcribe", "language": "en"}
        )
        logger.info("Whisper ASR pipeline initialised.")
        logger.info("Starting transcription...")

        # Perform transcription
        result = asr_pipeline(wav_file_path)
        transcription = result.get("text", "")
        logger.info("Transcription completed successfully.")

        yield transcription, log_buffer.getvalue()
    except Exception as e:
        err_msg = f"Error during transcription: {str(e)}"
        logger.error(err_msg)
        yield err_msg, log_buffer.getvalue()