Class: Transformers::PipelineDataset

Inherits:
Torch::Utils::Data::Dataset
  • Object
show all
Defined in:
lib/transformers/pipelines/pt_utils.rb

Instance Method Summary collapse

Constructor Details

#initialize(dataset, process, params) ⇒ PipelineDataset

Returns a new instance of PipelineDataset.



3
4
5
6
7
 
# File 'lib/transformers/pipelines/pt_utils.rb', line 3

def initialize(dataset, process, params)
  @dataset = dataset
  @process = process
  @params = params
end

Instance Method Details

#[](i) ⇒ Object 



13
14
15
16
17
 
# File 'lib/transformers/pipelines/pt_utils.rb', line 13

def [](i)
  item = @dataset[i]
  processed = @process.(item, **@params)
  processed
end

#size ⇒ Object 



9
10
11
 
# File 'lib/transformers/pipelines/pt_utils.rb', line 9

def size
  @dataset.size
end