Class: Transformers::PipelineRegistry

Inherits:
Object
  • Object
show all
Defined in:
lib/transformers/pipelines/base.rb

Instance Method Summary collapse

Constructor Details

#initialize(supported_tasks:, task_aliases:) ⇒ PipelineRegistry

Returns a new instance of PipelineRegistry.



279
280
281
282
 
# File 'lib/transformers/pipelines/base.rb', line 279

def initialize(supported_tasks:, task_aliases:)
  @supported_tasks = supported_tasks
  @task_aliases = task_aliases
end

Instance Method Details

#check_task(task) ⇒ Object

Raises:

  • (KeyError) 


289
290
291
292
293
294
295
296
297
298
299
 
# File 'lib/transformers/pipelines/base.rb', line 289

def check_task(task)
  if @task_aliases[task]
    task = @task_aliases[task]
  end
  if @supported_tasks[task]
    targeted_task = @supported_tasks[task]
    return task, targeted_task, nil
  end

  raise KeyError, "Unknown task #{task}, available tasks are #{get_supported_tasks}"
end

#get_supported_tasks ⇒ Object 



284
285
286
287
 
# File 'lib/transformers/pipelines/base.rb', line 284

def get_supported_tasks
  supported_task = @supported_tasks.keys + @task_aliases.keys
  supported_task.sort
end