llama-index

1
import re
2
from pathlib import Path
3
from typing import Dict, List, Optional
4

5
from llama_index.legacy.readers.base import BaseReader
6
from llama_index.legacy.schema import Document
7

8

9
class IPYNBReader(BaseReader):
10
    """Image parser."""
11

12
    def __init__(
13
        self,
14
        parser_config: Optional[Dict] = None,
15
        concatenate: bool = False,
16
    ):
17
        """Init params."""
18
        self._parser_config = parser_config
19
        self._concatenate = concatenate
20

21
    def load_data(
22
        self, file: Path, extra_info: Optional[Dict] = None
23
    ) -> List[Document]:
24
        """Parse file."""
25
        if file.name.endswith(".ipynb"):
26
            try:
27
                import nbconvert
28
            except ImportError:
29
                raise ImportError("Please install nbconvert 'pip install nbconvert' ")
30
        string = nbconvert.exporters.ScriptExporter().from_file(file)[0]
31
        # split each In[] cell into a separate string
32
        splits = re.split(r"In\[\d+\]:", string)
33
        # remove the first element, which is empty
34
        splits.pop(0)
35

36
        if self._concatenate:
37
            docs = [Document(text="\n\n".join(splits), metadata=extra_info or {})]
38
        else:
39
            docs = [Document(text=s, metadata=extra_info or {}) for s in splits]
40
        return docs
41