Refactor

Author: elcuervo

README.md

@@ -29,9 +29,19 @@ text = "Apple CEO Tim Cook announced iPhone 15 in Cupertino yesterday."
 labels = ["company", "person", "product", "location"]
 
 model = Gliner[labels]
-pp model[text]
+entities = model[text]
 
-# => {"company"=>["Apple"], "person"=>["Tim Cook"], "product"=>["iPhone 15"], "location"=>["Cupertino"]}
+pp entities["person"]
+# => [#<data Gliner::Entity ...>]
+
+entities["person"].first.text
+# => "Tim Cook"
+
+entities["person"].first.confidence
+# => 92.4
+
+entities["person"].first.offsets
+# => [10, 18]
 ```
 
 You can also pass per-entity configs:
@@ -43,9 +53,13 @@ labels = {
 }
 
 model = Gliner[labels]
-pp model["Email John Doe at john@example.com.", threshold: 0.5]
+entities = model["Email John Doe at john@example.com.", threshold: 0.5]
 
-# => {"email"=>["john@example.com"], "person"=>"John Doe"}
+entities["person"].text
+# => "John Doe"
+
+entities["email"].map(&:text)
+# => ["john@example.com"]
 ```
 
 ### Classification
@@ -59,7 +73,30 @@ result = model["This laptop has amazing performance but terrible battery life!"]
 
 pp result
 
-# => {"sentiment"=>"negative"}
+# => {"sentiment"=>#<data Gliner::Label ...>}
+
+result["sentiment"].label
+# => "negative"
+
+result["sentiment"].confidence
+# => 87.1
+```
+
+Multiple classification tasks:
+
+```ruby
+text = "Breaking: Tech giant announces major layoffs amid market downturn"
+
+tasks = {
+  "sentiment" => %w[positive negative neutral],
+  "urgency" => %w[high medium low],
+  "category" => { "labels" => %w[tech finance politics sports], "multi_label" => false }
+}
+
+results = Gliner.classify[tasks][text]
+
+results.transform_values { |value| value.label }
+# => {"sentiment"=>"negative", "urgency"=>"high", "category"=>"tech"}
 ```
 
 ### Structured extraction
@@ -77,10 +114,21 @@ structure = {
 }
 
 result = Gliner[structure][text]
+product = result.fetch("product").first
 
 pp result
 
-# => {"product"=>[{"name"=>"iPhone 15 Pro Max", "storage"=>"256GB", "processor"=>"A17 Pro", "price"=>"1199"}]}
+product["name"].text
+# => "iPhone 15 Pro Max"
+
+product["storage"].text
+# => "256GB"
+
+product["processor"].text
+# => "A17 Pro"
+
+product["price"].text
+# => "$1199"
 ```
 
 Choices can be included in field specs:

SPEC.md

@@ -5,9 +5,12 @@ text = "Apple CEO Tim Cook announced iPhone 15 in Cupertino yesterday."
 
 entities = ["company", "person", "product", "location"]
 model = Gliner[entities]
-model[text]
+entities = model[text]
 
-# => {company"=>["Apple"], "person"=>["Tim Cook"], "product"=>["iPhone 15"], "location"=>["Cupertino"]}
+# => {"company"=>[#<data Gliner::Entity ...>], "person"=>[#<data Gliner::Entity ...>], ...}
+
+entities["person"].first.text
+# => "Tim Cook"
 ```
 
 # Per entity config
@@ -21,9 +24,13 @@ entities = {
 }
 
 model = Gliner[entities]
-model[text]
+entities = model[text]
+
+entities["person"].text
+# => "John Doe"
 
-# => {company"=>["Apple"], "person"=>["Tim Cook"], "product"=>["iPhone 15"], "location"=>["Cupertino"]}
+entities["email"].map(&:text)
+# => ["john@example.com"]
 ```
 
 # Classfication
@@ -32,9 +39,15 @@ model[text]
 text = "This laptop has amazing performance but terrible battery life!",
 concept = { "sentiment" => %w[positive negative neutral] }
 model = Gliner.classify[concept]
-model[text]
+result = model[text]
+
+# => {"sentiment"=>#<data Gliner::Label ...>}
+
+result["sentiment"].label
+# => "negative"
 
-# => {"sentiment"=>"negative"}
+result["sentiment"].confidence
+# => 87.1
 ```
 
 # Structured
@@ -52,7 +65,18 @@ structure = {
 }
 
 model = Gliner[structure]
-model[text]
+result = model[text]
+product = result.fetch("product").first
+
+product["name"].text
+# => "iPhone 15 Pro Max"
+
+product["storage"].text
+# => "256GB"
+
+product["processor"].text
+# => "A17 Pro chip"
 
-# => {"product"=>[{"name"=>"iPhone 15 Pro Max", "storage"=>"256GB", "processor"=>"A17 Pro chip", "price"=>"$1199"}]}
+product["price"].text
+# => "$1199"
 ```

lib/gliner.rb

@@ -31,7 +31,7 @@ module Gliner
     :text_len
   )
 
-  Entity = Data.define(:index, :offsets, :text, :name, :probability) do
+  Entity = Data.define(:index, :offsets, :text, :name, :confidence) do
     def to_s = text.to_s
     def to_str = text.to_s
   end
@@ -47,17 +47,6 @@ def overlaps?(other)
     end
   end
 
-  FormatOptions = Data.define(:include_confidence, :include_spans) do
-    def self.from(input)
-      return input if input.is_a?(FormatOptions)
-
-      new(
-        include_confidence: input.fetch(:include_confidence, false),
-        include_spans: input.fetch(:include_spans, false)
-      )
-    end
-  end
-
   class << self
     attr_writer :model, :config
 

lib/gliner/classifier.rb

@@ -47,7 +47,9 @@ def sorted_label_scores(scores, labels)
     def format_multi_label(label_scores, cls_threshold, include_confidence)
       chosen = labels_above_threshold(label_scores, cls_threshold)
 
-      chosen.map { |label, score| format_label(label, score, include_confidence) }
+      chosen
+        .sort_by { |(_label, score)| -score }
+        .map { |label, score| format_label(label, score, include_confidence) }
     end
 
     def labels_above_threshold(label_scores, threshold)

lib/gliner/span_extractor.rb

@@ -109,7 +109,7 @@ def format_span(span, _opts, label:, index:)
         offsets: [span.start, span.end],
         text: span.text,
         name: label&.to_s,
-        probability: span.score * 100.0
+        confidence: span.score * 100.0
       )
     end
 

lib/gliner/structured_extractor.rb

@@ -2,6 +2,8 @@
 
 module Gliner
   Structure = Data.define(:fields) do
+    include Enumerable
+
     def [](key) = fields[key]
     def fetch(key, *args, &block) = fields.fetch(key, *args, &block)
     def to_h = fields
@@ -42,15 +44,14 @@ def filter_spans_by_choices(spans, choices)
     end
 
     def build_structure_instances(parsed_fields, spans_by_label, opts)
-      format_opts = FormatOptions.from(opts)
       anchor_field = anchor_field_for(parsed_fields)
       return [Gliner::Structure.new(fields: {})] unless anchor_field
 
       anchors = spans_by_label.fetch(anchor_field[:name], [])
-      return [format_structure_object(parsed_fields, spans_by_label, format_opts)] if anchors.empty?
+      return [format_structure_object(parsed_fields, spans_by_label, opts)] if anchors.empty?
 
       instance_spans = build_instance_spans(anchors, spans_by_label)
-      format_instances(parsed_fields, instance_spans, format_opts)
+      format_instances(parsed_fields, instance_spans, opts)
     end
 
     def format_structure_object(parsed_fields, spans_by_label, _opts)

lib/gliner/tasks/json_extraction.rb

@@ -47,9 +47,7 @@ def labels(parsed)
       def process_output(logits, parsed, prepared, options)
         spans_by_label = extract_spans(logits, parsed, prepared, options)
         filtered_spans = @structured_extractor.apply_choice_filters(spans_by_label, parsed[:parsed_fields])
-        format_opts = FormatOptions.from(options)
-
-        @structured_extractor.build_structure_instances(parsed[:parsed_fields], filtered_spans, format_opts)
+        @structured_extractor.build_structure_instances(parsed[:parsed_fields], filtered_spans, options)
       end
 
       def execute_all(pipeline, text, structures_config, **options)

spec/gliner/model_spec.rb

@@ -22,7 +22,7 @@
           offsets: [0, 8],
           text: 'Tim Cook',
           name: 'person',
-          probability: 90.0
+          confidence: 90.0
         )
       )
     end
@@ -109,21 +109,21 @@
             offsets: [0, 5],
             text: 'Jan 5',
             name: 'date',
-            probability: 90.0
+            confidence: 90.0
           ),
           'merchant' => Gliner::Entity.new(
             index: 0,
             offsets: [7, 16],
             text: 'Starbucks',
             name: 'merchant',
-            probability: 90.0
+            confidence: 90.0
           ),
           'amount' => Gliner::Entity.new(
             index: 0,
             offsets: [17, 22],
             text: '$5.50',
             name: 'amount',
-            probability: 90.0
+            confidence: 90.0
           )
         })
       )
@@ -134,21 +134,21 @@
             offsets: [40, 45],
             text: 'Jan 6',
             name: 'date',
-            probability: 80.0
+            confidence: 80.0
           ),
           'merchant' => Gliner::Entity.new(
             index: 0,
             offsets: [47, 53],
             text: 'Amazon',
             name: 'merchant',
-            probability: 70.0
+            confidence: 70.0
           ),
           'amount' => Gliner::Entity.new(
             index: 0,
             offsets: [54, 61],
             text: '$156.99',
             name: 'amount',
-            probability: 80.0
+            confidence: 80.0
           )
         })
       )