Tools/transcriber.py

from smolagents import tool
import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline

@tool
def transcriber(audio_file_path:str) -> str:
    """
    This tool transcribes an audio file and returns the generated transcription.

    Args:
        audio_file_path: Path of a local audio file that needs to be transcribed.
    """

    try:
        device = "cuda:0" if torch.cuda.is_available() else "cpu"
        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

        model_id = "openai/whisper-small"

        model = AutoModelForSpeechSeq2Seq.from_pretrained(
            model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
        )
        model.to(device)

        processor = AutoProcessor.from_pretrained(model_id)

        pipe = pipeline(
            "automatic-speech-recognition",
            model=model,
            tokenizer=processor.tokenizer,
            feature_extractor=processor.feature_extractor,
            torch_dtype=torch_dtype,
            device=device,
            return_timestamps=True
        )

        result = pipe(audio_file_path)

        import gc
        
        # After inference
        del pipe
        del model
        del processor
        gc.collect()                 # Force Python garbage collection
        torch.cuda.empty_cache()     # Clear cached memory
        
        return result["text"]
    except Exception as e:
        return f'error occured: {e}'