from smolagents import Tool
import whisper

class AudioTranscriberTool(Tool):
    name = "audio_transcriber"
    description = ("""
        This is a tool that get the transcription of the audio file in the form of text. 
        Supported file extensions: .mp3, .wav, .flac, .ogg, .m4a.
    """)
    inputs = {
        "file_path": {
            "type": "string",
            "description": "The file path to the audio file.",
        }
    }
    output_type = "string"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model = whisper.load_model("base")

    def forward(self, file_path: str) -> str:
        result = self.model.transcribe(file_path)
        return result['text']