File size: 931 Bytes
5445ab9
d2a63cc
 
 
5445ab9
 
 
 
d2a63cc
e4316f1
31086ae
 
 
5445ab9
31086ae
 
 
 
 
5445ab9
 
31086ae
 
0bd1b0f
31086ae
 
 
 
 
 
0bd1b0f
31086ae
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
from typing import TYPE_CHECKING, List, Union

from graphgen.bases.base_reader import BaseReader

if TYPE_CHECKING:
    import ray
    from ray.data import Dataset


class TXTReader(BaseReader):
    def read(
        self,
        input_path: Union[str, List[str]],
    ) -> "Dataset":
        """
        Read text files from the specified input path.
        :param input_path: Path to the input text file or list of text files.
        :return: Ray Dataset containing the read text data.
        """
        import ray

        docs_ds = ray.data.read_binary_files(
            input_path,
            include_paths=True,
        )

        docs_ds = docs_ds.map(
            lambda row: {
                "type": "text",
                self.text_column: row["bytes"].decode("utf-8"),
                "path": row["path"],
            }
        )

        docs_ds = docs_ds.filter(self._should_keep_item)
        return docs_ds