Spaces:
Running
Running
File size: 799 Bytes
31086ae d2a63cc 31086ae d2a63cc e4316f1 52419fe 31086ae d2a63cc 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 |
from typing import List, Union
import ray
from ray.data import Dataset
from graphgen.bases.base_reader import BaseReader
class CSVReader(BaseReader):
"""
Reader for CSV files.
Columns:
- type: The type of the document (e.g., "text", "image", etc.)
- if type is "text", "content" column must be present.
"""
def read(self, input_path: Union[str, List[str]]) -> Dataset:
"""
Read CSV files and return Ray Dataset.
:param input_path: Path to CSV file or list of CSV files.
:return: Ray Dataset containing validated and filtered data.
"""
ds = ray.data.read_csv(input_path)
ds = ds.map_batches(self._validate_batch, batch_format="pandas")
ds = ds.filter(self._should_keep_item)
return ds
|