Elron commited on
Commit
8fb5471
1 Parent(s): d423f18

Upload load.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. load.py +20 -0
load.py ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from .operator import StreamSource
2
+ from datasets import DatasetDict
3
+ from typing import Union
4
+
5
+ from .catalog import LocalCatalog
6
+ from .artifact import Artifact
7
+
8
+
9
+ def load_stream(source_name_or_path: str) -> StreamSource:
10
+ if Artifact.is_artifact_file(source_name_or_path):
11
+ return Artifact.load(source_name_or_path)
12
+ else:
13
+ return LocalCatalog().load(source_name_or_path)
14
+
15
+
16
+ def load_dataset(source: Union[StreamSource, str]) -> DatasetDict:
17
+ assert isinstance(source, (StreamSource, str)), "source must be a StreamSource or a string"
18
+ if isinstance(source, str):
19
+ source = load_stream(source)
20
+ return source().to_dataset()