More generic core DataLoader

epfml · Feb 3, 2024 · cb67111 · cb67111
1 parent 2d0ab0d
commit cb67111
Showing 1 changed file with 17 additions and 7 deletions.
diff --git a/discojs/discojs-core/src/dataset/data_loader/data_loader.ts b/discojs/discojs-core/src/dataset/data_loader/data_loader.ts
@@ -1,15 +1,25 @@
-import { Task } from '../..'
+import { Task } from '../../'
 import { Dataset } from '../dataset'
 import { Data, DataSplit } from '../data'
 
-export interface DataConfig { features?: string[], labels?: string[], shuffle?: boolean, validationSplit?: number, inference?: boolean }
+export interface DataConfig {
+    features?: string[]
+    labels?: string[]
+    shuffle?: boolean
+    validationSplit?: number
+    inference?: boolean
+}
 
-export abstract class DataLoader<Source> {
-  constructor (protected task: Task) {}
+export abstract class DataLoader<
+    Source,
+    Sources = Source | Source[],
+    Config extends DataConfig = DataConfig
+> {
+    constructor(protected task: Task) {}
 
-  abstract createData (dataset: Dataset, size?: number): Promise<Data>
+    abstract createData(dataset: Dataset, size?: number): Promise<Data>
 
-  abstract load (source: Source, config: DataConfig): Promise<Dataset>
+    abstract load(source: Source, config: Config): Promise<Dataset>
 
-  abstract loadAll (sources: Source[], config: DataConfig): Promise<DataSplit>
+    abstract loadAll(sources: Sources, config: Config): Promise<DataSplit>
 }