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.



316
317
318
319
 
# File 'lib/transformers/pipelines/base.rb', line 316

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

Instance Method Details

#check_task(task) ⇒ Object

Raises:

  • (KeyError) 


326
327
328
329
330
331
332
333
334
335
336
 
# File 'lib/transformers/pipelines/base.rb', line 326

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 



321
322
323
324
 
# File 'lib/transformers/pipelines/base.rb', line 321

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