mirror of
https://github.com/run-llama/llama-hub.git
synced 2025-08-15 20:21:14 +00:00
65 lines
1.8 KiB
Python
65 lines
1.8 KiB
Python
"""Audio Transcriber.
|
|
|
|
A transcriber for the audio of mp3, mp4 files.
|
|
|
|
"""
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional, List, cast
|
|
|
|
from gpt_index.readers.base import BaseReader
|
|
from gpt_index.readers.schema.base import Document
|
|
|
|
|
|
class AudioTranscriber(BaseReader):
|
|
"""Audio parser.
|
|
|
|
Extract text from transcript of video/audio files using OpenAI Whisper.
|
|
|
|
"""
|
|
|
|
def __init__(self, *args: Any, model_version: str = "base", **kwargs: Any) -> None:
|
|
"""Init params."""
|
|
super().__init__(*args, **kwargs)
|
|
self._model_version = model_version
|
|
|
|
try:
|
|
import whisper
|
|
except ImportError:
|
|
raise ValueError(
|
|
"Please install OpenAI whisper model "
|
|
"'pip install git+https://github.com/openai/whisper.git' "
|
|
"to use the model"
|
|
)
|
|
|
|
model = whisper.load_model(self._model_version)
|
|
|
|
self.parser_config = {"model": model}
|
|
|
|
def load_data(
|
|
self, file: Path, extra_info: Optional[Dict] = None
|
|
) -> List[Document]:
|
|
"""Parse file."""
|
|
import whisper
|
|
|
|
if file.name.endswith("mp4"):
|
|
try:
|
|
from pydub import AudioSegment # noqa: F401
|
|
except ImportError:
|
|
raise ValueError("Please install pydub 'pip install pydub' ")
|
|
# open file
|
|
video = AudioSegment.from_file(file, format="mp4")
|
|
|
|
# Extract audio from video
|
|
audio = video.split_to_mono()[0]
|
|
|
|
file_str = str(file)[:-4] + ".mp3"
|
|
# export file
|
|
audio.export(file_str, format="mp3")
|
|
|
|
model = cast(whisper.Whisper, self.parser_config["model"])
|
|
result = model.transcribe(str(file))
|
|
|
|
transcript = result["text"]
|
|
|
|
return [Document(transcript, extra_info=extra_info)]
|