mirror of
https://github.com/wassname/pytorch-ts.git
synced 2026-08-20 12:40:22 +08:00
134 lines
3.8 KiB
Python
134 lines
3.8 KiB
Python
# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License").
|
|
# You may not use this file except in compliance with the License.
|
|
# A copy of the License is located at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# or in the "license" file accompanying this file. This file is distributed
|
|
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
|
|
# express or implied. See the License for the specific language governing
|
|
# permissions and limitations under the License.
|
|
|
|
|
|
import functools
|
|
import glob
|
|
import random
|
|
from pathlib import Path
|
|
from typing import Iterator, List
|
|
from typing import NamedTuple
|
|
|
|
import rapidjson as json
|
|
|
|
from .common import Dataset, DataEntry, SourceContext
|
|
from .process import ProcessDataEntry
|
|
|
|
|
|
def load(file_obj):
|
|
for line in file_obj:
|
|
yield json.loads(line)
|
|
|
|
|
|
class Span(NamedTuple):
|
|
path: Path
|
|
line: int
|
|
|
|
|
|
class Line(NamedTuple):
|
|
content: object
|
|
span: Span
|
|
|
|
|
|
class JsonLinesFile:
|
|
"""
|
|
An iterable type that draws from a JSON Lines file.
|
|
|
|
Parameters
|
|
----------
|
|
path
|
|
Path of the file to load data from. This should be a valid
|
|
JSON Lines file.
|
|
"""
|
|
|
|
def __init__(self, path: Path, shuffle: bool = True) -> None:
|
|
self.path = path
|
|
self.shuffle = shuffle
|
|
|
|
def __iter__(self):
|
|
with open(self.path) as jsonl_file:
|
|
lines = jsonl_file.read().splitlines()
|
|
if self.shuffle:
|
|
random.shuffle(lines)
|
|
|
|
for line_number, raw in enumerate(lines, start=1):
|
|
span = Span(path=self.path, line=line_number)
|
|
try:
|
|
yield Line(json.loads(raw), span=span)
|
|
except ValueError:
|
|
raise Exception(f"Could not read json line {line_number}, {raw}")
|
|
|
|
def __len__(self):
|
|
# 1MB
|
|
BUF_SIZE = 1024 ** 2
|
|
|
|
with open(self.path) as file_obj:
|
|
read_chunk = functools.partial(file_obj.read, BUF_SIZE)
|
|
return sum(chunk.count("\n") for chunk in iter(read_chunk, ""))
|
|
|
|
|
|
class FileDataset(Dataset):
|
|
"""
|
|
Dataset that loads JSON Lines files contained in a path.
|
|
|
|
Parameters
|
|
----------
|
|
path
|
|
Return list of path names that match path. Each file is considered
|
|
and should be valid. A valid line in a file can be for
|
|
instance: {"start": "2014-09-07", "target": [0.1, 0.2]}.
|
|
freq
|
|
Frequency of the observation in the time series.
|
|
Must be a valid Pandas frequency.
|
|
one_dim_target
|
|
Whether to accept only univariate target time series.
|
|
shuffle
|
|
Whether to shuffle the time series when making the batches
|
|
"""
|
|
|
|
def __init__(
|
|
self, path: Path, freq: str, one_dim_target: bool = True, shuffle: bool = False
|
|
) -> None:
|
|
self.shuffle = shuffle
|
|
self.path = path
|
|
self.process = ProcessDataEntry(freq, one_dim_target=one_dim_target)
|
|
if not self.files():
|
|
raise OSError(f"no valid file found via {path}")
|
|
|
|
def __iter__(self) -> Iterator[DataEntry]:
|
|
for path in self.files():
|
|
for line in JsonLinesFile(path, self.shuffle):
|
|
data = self.process(line.content)
|
|
data["source"] = SourceContext(
|
|
source=line.span.path, row=line.span.line
|
|
)
|
|
yield data
|
|
|
|
def __len__(self):
|
|
return sum([len(JsonLinesFile(path)) for path in self.files()])
|
|
|
|
def files(self) -> List[Path]:
|
|
"""
|
|
List the files that compose the dataset.
|
|
|
|
Returns
|
|
-------
|
|
List[Path]
|
|
List of the paths of all files composing the dataset.
|
|
"""
|
|
files = glob.glob(str(self.path))
|
|
if self.shuffle:
|
|
random.shuffle(files)
|
|
return files
|
|
|