Class: Gliner::Tasks::EntityExtraction

Inherits:
Gliner::Task show all
Defined in:
lib/gliner/tasks/entity_extraction.rb

Instance Attribute Summary

Attributes inherited from Gliner::Task

#config_parser, #inference, #input_builder

Instance Method Summary collapse

Methods inherited from Gliner::Task

#needs_cls_logits?, #normalize_text?

Constructor Details

#initialize(config_parser:, inference:, input_builder:, span_extractor:) ⇒ EntityExtraction

Returns a new instance of EntityExtraction.



6
7
8
9
 
# File 'lib/gliner/tasks/entity_extraction.rb', line 6

def initialize(config_parser:, inference:, input_builder:, span_extractor:)
  super(config_parser: config_parser, inference: inference, input_builder: input_builder)
  @span_extractor = span_extractor
end

Instance Method Details

#build_prompt(parsed) ⇒ Object 



23
24
25
 
# File 'lib/gliner/tasks/entity_extraction.rb', line 23

def build_prompt(parsed)
  config_parser.build_prompt('entities', parsed[:descriptions])
end

#label_prefix ⇒ Object 



19
20
21
 
# File 'lib/gliner/tasks/entity_extraction.rb', line 19

def label_prefix
  '[E]'
end

#labels(parsed) ⇒ Object 



27
28
29
 
# File 'lib/gliner/tasks/entity_extraction.rb', line 27

def labels(parsed)
  parsed[:labels]
end

#parse_config(input) ⇒ Object 



11
12
13
 
# File 'lib/gliner/tasks/entity_extraction.rb', line 11

def parse_config(input)
  config_parser.parse_entity_types(input)
end

#process_output(logits, parsed, prepared, options) ⇒ Object 



31
32
33
34
35
36
37
38
39
 
# File 'lib/gliner/tasks/entity_extraction.rb', line 31

def process_output(logits, parsed, prepared, options)
  threshold = options.fetch(:threshold, Gliner.config.threshold)
  format_opts = FormatOptions.from(options)
  label_positions = options[:label_positions] || inference.label_positions_for(prepared.word_ids, parsed[:labels].length)

  spans_by_label = extract_spans(logits, parsed, prepared, label_positions, threshold)

  { 'entities' => format_entities(parsed, spans_by_label, format_opts) }
end

#task_type ⇒ Object 



15
16
17
 
# File 'lib/gliner/tasks/entity_extraction.rb', line 15

def task_type
  Inference::TASK_TYPE_ENTITIES
end