llama-index

1
"""Video audio parser.
2

3
Contains parsers for mp3, mp4 files.
4

5
"""
6

7
from pathlib import Path
8
from typing import Any, Dict, List, Optional, cast
9

10
from llama_index.legacy.readers.base import BaseReader
11
from llama_index.legacy.schema import Document
12

13

14
class VideoAudioReader(BaseReader):
15
    """Video audio parser.
16

17
    Extract text from transcript of video/audio files.
18

19
    """
20

21
    def __init__(self, *args: Any, model_version: str = "base", **kwargs: Any) -> None:
22
        """Init parser."""
23
        super().__init__(*args, **kwargs)
24
        self._model_version = model_version
25

26
        try:
27
            import whisper
28
        except ImportError:
29
            raise ImportError(
30
                "Please install OpenAI whisper model "
31
                "'pip install git+https://github.com/openai/whisper.git' "
32
                "to use the model"
33
            )
34

35
        model = whisper.load_model(self._model_version)
36

37
        self.parser_config = {"model": model}
38

39
    def load_data(
40
        self, file: Path, extra_info: Optional[Dict] = None
41
    ) -> List[Document]:
42
        """Parse file."""
43
        import whisper
44

45
        if file.name.endswith("mp4"):
46
            try:
47
                from pydub import AudioSegment
48
            except ImportError:
49
                raise ImportError("Please install pydub 'pip install pydub' ")
50
            # open file
51
            video = AudioSegment.from_file(file, format="mp4")
52

53
            # Extract audio from video
54
            audio = video.split_to_mono()[0]
55

56
            file_str = str(file)[:-4] + ".mp3"
57
            # export file
58
            audio.export(file_str, format="mp3")
59

60
        model = cast(whisper.Whisper, self.parser_config["model"])
61
        result = model.transcribe(str(file))
62

63
        transcript = result["text"]
64

65
        return [Document(text=transcript, metadata=extra_info or {})]
66