Class: Cohere::Transcribe::ASR::NativeSession

Inherits:
Object
  • Object
show all
Defined in:
lib/cohere/transcribe/asr/native.rb

Overview

One retained native Cohere model session.

Constant Summary collapse

MAX_NATIVE_SAMPLE_COUNT = 
2_147_483_647
NATIVE_CANCELLED_KIND = 
5
CANCELLATION_JOIN_INTERVAL = 
0.01
NATIVE_BATCH_STAT_FIELDS = 
%i[
  abi_version field_count total_us feature_wall_us feature_worker_us mel_pack_us
  encoder_graph_build_us encoder_graph_alloc_us encoder_input_us encoder_compute_us
  encoder_readback_us encoder_repack_us decoder_total_us decoder_cross_kv_us
  decoder_reserve_us decoder_decode_us render_us decoder_calls generation_steps
  encoder_microbatches token_id_readback_bytes
].freeze
NATIVE_BATCH_STAT_ABI_VERSION = 
1
NATIVE_FAILURE_KINDS = 
{
  1 => :fatal,
  2 => :oom,
  3 => :fatal,
  4 => :error
}.freeze

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(model_path, options, threads: nil, library: NativeLibrary.load) ⇒ NativeSession

Returns a new instance of NativeSession.



313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
 
# File 'lib/cohere/transcribe/asr/native.rb', line 313

def initialize(model_path, options, threads: nil, library: NativeLibrary.load)
  initialized = false
  @library = library
  @model_path = Pathname.new(model_path).expand_path.freeze
  raise TranscriptionRuntimeError, "Native Dense model does not exist: #{@model_path}" unless @model_path.file?

  @mutex = Mutex.new
  @closed = false
  @session_ownership = Internal::SessionOwnership.new(
    close: NativeSessionCloser.new(@library),
    installed_error: "Native session ownership is already installed"
  )
  ObjectSpace.define_finalizer(
    self,
    Internal::SessionOwnership.finalizer_for(@session_ownership)
  )
  thread_count = normalize_threads(threads)
  # An asynchronous exception may become deliverable immediately after
  # the foreign open returns. Defer it until the pointer belongs to the
  # independent cleanup state, then let constructor teardown close it.
  Thread.handle_interrupt(Object => :never) do
    session = @library.open_session(
      model_path: @model_path,
      device: options.device.to_s,
      threads: thread_count
    )
    @session_ownership.install(session)
    @session = session
  end
  @default_max_new_tokens = Integer(options.max_new_tokens)
  @library.call(:session_set_max_new_tokens, @session, @default_max_new_tokens)
  @library.call(:session_set_beam_size, @session, 1)
  @library.call(
    :session_set_repetition_loop_guard,
    @session,
    options.stop_repetition_loops ? 1 : 0
  )
  @backend = @library.string(@library.call(:session_backend, @session)).freeze
  unless @backend == "cohere"
    close
    raise TranscriptionRuntimeError,
          "Native runtime selected #{@backend.inspect}, expected the Cohere Dense backend"
  end
  @compute_backend = @library.string(
    @library.call(:session_compute_backend, @session)
  ).freeze
  @device = canonical_device(@compute_backend)
  unless @device == options.device.to_s
    close
    raise TranscriptionRuntimeError,
          "Native runtime selected #{@compute_backend.inspect} (#{@device}) " \
          "after #{options.device.inspect} was resolved"
  end
  @batch_capacity = Integer(@library.call(:session_batch_capacity, @session))
  unless @batch_capacity.positive?
    close
    raise TranscriptionRuntimeError, "Native runtime reported an invalid Cohere batch capacity"
  end
  initialized = true
ensure
  unless initialized
    Thread.handle_interrupt(Object => :never) do
      close if defined?(@session_ownership)
    rescue Exception # rubocop:disable Lint/RescueException -- preserve constructor termination
      nil
    end
  end
end

Instance Attribute Details

#backend ⇒ Object (readonly)

Returns the value of attribute backend.



311
312
313
 
# File 'lib/cohere/transcribe/asr/native.rb', line 311

def backend
  @backend
end

#batch_capacity ⇒ Object (readonly)

Returns the value of attribute batch_capacity.



311
312
313
 
# File 'lib/cohere/transcribe/asr/native.rb', line 311

def batch_capacity
  @batch_capacity
end

#compute_backend ⇒ Object (readonly)

Returns the value of attribute compute_backend.



311
312
313
 
# File 'lib/cohere/transcribe/asr/native.rb', line 311

def compute_backend
  @compute_backend
end

#device ⇒ Object (readonly)

Returns the value of attribute device.



311
312
313
 
# File 'lib/cohere/transcribe/asr/native.rb', line 311

def device
  @device
end

#last_batch_metrics ⇒ Object (readonly)

Returns the value of attribute last_batch_metrics.



311
312
313
 
# File 'lib/cohere/transcribe/asr/native.rb', line 311

def last_batch_metrics
  @last_batch_metrics
end

#model_path ⇒ Object (readonly)

Returns the value of attribute model_path.



311
312
313
 
# File 'lib/cohere/transcribe/asr/native.rb', line 311

def model_path
  @model_path
end

Instance Method Details

#close ⇒ Object 



485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
 
# File 'lib/cohere/transcribe/asr/native.rb', line 485

def close
  @mutex ||= Mutex.new
  @mutex.synchronize do
    return if @closed

    @session = nil
    @closed = true
    begin
      @session_ownership&.close
    ensure
      ObjectSpace.undefine_finalizer(self)
    end
  end
  nil
end

#closed? ⇒ Boolean

Returns:

  • (Boolean) 


501
502
503
 
# File 'lib/cohere/transcribe/asr/native.rb', line 501

def closed?
  @closed
end

#memory ⇒ Object

Current free/total memory for the selected ggml device, used by the adaptive batch controller. Returns nil when a backend has no memory telemetry rather than guessing from process RSS.



508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
 
# File 'lib/cohere/transcribe/asr/native.rb', line 508

def memory
  @mutex.synchronize do
    ensure_open!
    free_bytes = [0].pack("Q")
    total_bytes = [0].pack("Q")
    status = @library.call(
      :session_memory,
      @session,
      Fiddle::Pointer[free_bytes],
      Fiddle::Pointer[total_bytes]
    )
    return nil unless status.zero?

    [free_bytes.unpack1("Q"), total_bytes.unpack1("Q")].freeze
  end
end

#transcribe(samples, language:, offset: 0.0, max_new_tokens: nil) ⇒ Object 



382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
 
# File 'lib/cohere/transcribe/asr/native.rb', line 382

def transcribe(samples, language:, offset: 0.0, max_new_tokens: nil)
  @mutex.synchronize do
    ensure_open!
    @last_batch_metrics = nil
    binary, sample_count = float_samples(samples)
    raise ArgumentError, "Audio segment must not be empty" if sample_count.zero?
    raise TranscriptionRuntimeError, "Audio segment is too large for the native ABI" \
      if sample_count > MAX_NATIVE_SAMPLE_COUNT

    generation_limit = max_new_tokens.nil? ? @default_max_new_tokens : Integer(max_new_tokens)
    set_generation_limit!(generation_limit)

    run_native_inference do
      result = @library.call(
        :session_transcribe_lang,
        @session,
        Fiddle::Pointer[binary],
        sample_count,
        Fiddle::Pointer["#{language}\0"]
      )
      raise_native_failure!("Native Cohere inference returned no result") if @library.null_pointer?(result)

      begin
        materialize_result(result, offset: Float(offset))
      ensure
        @library.call(:session_result_free, result)
      end
    end
  end
end

#transcribe_batch(sample_batches, language:, offsets: nil, max_new_tokens: nil) ⇒ Object

True padded-encoder/ragged-decoder batching in the native Cohere runtime. The capacity is queried from the loaded adapter because its logical decoder batch may span multiple encoder microbatches.



416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
 
# File 'lib/cohere/transcribe/asr/native.rb', line 416

def transcribe_batch(sample_batches, language:, offsets: nil, max_new_tokens: nil)
  unless sample_batches.is_a?(Array) && sample_batches.length.between?(1, @batch_capacity)
    raise ArgumentError,
          "sample_batches must contain between 1 and #{@batch_capacity} audio rows"
  end

  offsets ||= Array.new(sample_batches.length, 0.0)
  unless offsets.is_a?(Array) && offsets.length == sample_batches.length
    raise ArgumentError, "offsets must contain one value per audio row"
  end

  @mutex.synchronize do
    ensure_open!
    @last_batch_metrics = nil
    buffers_and_counts = sample_batches.map { |samples| float_samples(samples) }
    empty_lane = buffers_and_counts.index { |_binary, count| count.zero? }
    raise ArgumentError, "Audio batch row #{empty_lane} must not be empty" if empty_lane

    oversized_lane = buffers_and_counts.index { |_binary, count| count > MAX_NATIVE_SAMPLE_COUNT }
    if oversized_lane
      raise TranscriptionRuntimeError,
            "Audio batch row #{oversized_lane} is too large for the native ABI"
    end

    generation_limit = max_new_tokens.nil? ? @default_max_new_tokens : Integer(max_new_tokens)
    set_generation_limit!(generation_limit)

    run_native_inference do
      pointers = buffers_and_counts.map { |binary, _count| Fiddle::Pointer[binary] }
      pointer_table = pointers.map(&:to_i).pack("J*")
      counts = buffers_and_counts.map(&:last).pack("i!*")
      batch_result = @library.call(
        :session_transcribe_batch_lang,
        @session,
        Fiddle::Pointer[pointer_table],
        Fiddle::Pointer[counts],
        sample_batches.length,
        Fiddle::Pointer["#{language}\0"]
      )
      raise_native_failure!("Native Cohere batch inference returned no result") \
        if @library.null_pointer?(batch_result)

      begin
        @last_batch_metrics = materialize_batch_metrics(batch_result)
        count = @library.call(:session_batch_result_count, batch_result)
        unless count == sample_batches.length
          raise ExecutionError.new(
            "Native Cohere batch returned #{count} rows for #{sample_batches.length} inputs",
            failure_kind: :fatal
          )
        end
        Array.new(count) do |lane|
          result = @library.call(:session_batch_result_at, batch_result, lane)
          if @library.null_pointer?(result)
            raise ExecutionError.new(
              "Native Cohere batch result row #{lane} is missing",
              failure_kind: :fatal
            )
          end

          materialize_result(result, offset: Float(offsets.fetch(lane)))
        end.freeze
      ensure
        @library.call(:session_batch_result_free, batch_result)
      end
    end
  end
end