Class: Cohere::Transcribe::Alignment::Aligner

Inherits:
Object
  • Object
show all
Defined in:
lib/cohere/transcribe/alignment/aligner.rb

Overview

Full-file MMS emissions plus per-ASR-segment CTC Viterbi alignment. The 30 s windows, 2 s context, 20 ms stride, wildcard column, and uniform per-segment recovery are the same as the Python reference.

Constant Summary collapse

SAMPLE_RATE = 
16_000
INPUTS_TO_LOGITS_RATIO = 
320
WINDOW_SECONDS = 
30
CONTEXT_SECONDS = 
2
WINDOW_SAMPLES = 
WINDOW_SECONDS * SAMPLE_RATE
CONTEXT_SAMPLES = 
CONTEXT_SECONDS * SAMPLE_RATE
INPUT_SAMPLES = 
WINDOW_SAMPLES + (2 * CONTEXT_SAMPLES)
MINIMUM_INPUT_SAMPLES = 
400
WINDOW_FRAMES = 
WINDOW_SAMPLES / INPUTS_TO_LOGITS_RATIO
CONTEXT_FRAMES = 
CONTEXT_SAMPLES / INPUTS_TO_LOGITS_RATIO
STRIDE_MS = 
INPUTS_TO_LOGITS_RATIO * 1_000.0 / SAMPLE_RATE
ISO3 = 
{ "ar" => "ara", "en" => "eng" }.freeze
VOCABULARY = 
{
  "<blank>" => 0, "<pad>" => 1, "</s>" => 2, "<unk>" => 3,
  "a" => 4, "i" => 5, "e" => 6, "n" => 7, "o" => 8, "u" => 9,
  "t" => 10, "s" => 11, "r" => 12, "m" => 13, "k" => 14,
  "l" => 15, "d" => 16, "g" => 17, "h" => 18, "y" => 19,
  "b" => 20, "p" => 21, "w" => 22, "c" => 23, "v" => 24,
  "j" => 25, "z" => 26, "f" => 27, "'" => 28, "q" => 29,
  "x" => 30
}.freeze
BLANK_ID = 
VOCABULARY.fetch("<blank>")

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(dtype: "fp32", device: "cpu", batch_size: 4, session: nil, **session_options) ⇒ Aligner

Returns a new instance of Aligner.

Raises:

  • (ArgumentError) 


252
253
254
255
256
257
258
259
260
 
# File 'lib/cohere/transcribe/alignment/aligner.rb', line 252

def initialize(dtype: "fp32", device: "cpu", batch_size: 4, session: nil, **session_options)
  raise ArgumentError, "aligner batch_size must be positive" unless batch_size.is_a?(Integer) && batch_size.positive?

  @session = session || Session.new(dtype: dtype, device: device, **session_options)
  @batch_size = batch_size
  @learned_batch_size = batch_size
  @emissions_seconds = 0.0
  @viterbi_seconds = 0.0
end

Instance Attribute Details

#batch_size ⇒ Object (readonly)

Returns the value of attribute batch_size.



250
251
252
 
# File 'lib/cohere/transcribe/alignment/aligner.rb', line 250

def batch_size
  @batch_size
end

#emissions_seconds ⇒ Object (readonly)

Returns the value of attribute emissions_seconds.



250
251
252
 
# File 'lib/cohere/transcribe/alignment/aligner.rb', line 250

def emissions_seconds
  @emissions_seconds
end

#session ⇒ Object (readonly)

Returns the value of attribute session.



250
251
252
 
# File 'lib/cohere/transcribe/alignment/aligner.rb', line 250

def session
  @session
end

#viterbi_seconds ⇒ Object (readonly)

Returns the value of attribute viterbi_seconds.



250
251
252
 
# File 'lib/cohere/transcribe/alignment/aligner.rb', line 250

def viterbi_seconds
  @viterbi_seconds
end

Instance Method Details

#align(audio, segment_times, segment_texts, language:) ⇒ Object 



279
280
281
282
283
284
285
286
287
 
# File 'lib/cohere/transcribe/alignment/aligner.rb', line 279

def align(audio, segment_times, segment_texts, language:)
  raise ArgumentError, "segment_times and segment_texts must have equal lengths" unless segment_times.length == segment_texts.length

  emissions, stride_ms = compute_emissions(audio)
  started = monotonic
  align_emissions(emissions, stride_ms, segment_times, segment_texts, language: language)
ensure
  @viterbi_seconds += monotonic - started if started
end

#align_emissions(emissions, stride_ms, segment_times, segment_texts, language:) ⇒ Object

Raises:



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
381
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
 
# File 'lib/cohere/transcribe/alignment/aligner.rb', line 349

def align_emissions(emissions, stride_ms, segment_times, segment_texts, language:)
  frame_count, emission_classes = emissions.shape
  star_id = emission_classes - 1
  raise TranscriptionRuntimeError, "MMS wildcard column collides with its tokenizer vocabulary" if VOCABULARY.value?(star_id)

  dictionary = VOCABULARY.merge("<star>" => star_id).freeze
  index_to_token = dictionary.invert.freeze
  iso_language = ISO3.fetch(language, language)
  words = []
  fallback_count = 0

  segment_times.zip(segment_texts).each_with_index do |((start_time, end_time), text), segment_index|
    text = PythonText.strip(text.to_s)
    next if text.empty?

    first_frame = [(start_time * 1_000 / stride_ms).round(half: :even), 0].max
    last_frame = [(end_time * 1_000 / stride_ms).round(half: :even), frame_count].min
    if last_frame - first_frame < 2
      fallback_count += 1
      words.concat(uniform_fallback(text, start_time, end_time, segment_index))
      next
    end

    begin
      tokens_starred, text_starred = Text.preprocess(text, iso_language)
      targets = tokens_starred.join(" ").split.filter_map { |token| dictionary[token] }
      raise ArgumentError, "Transcript produced no aligner vocabulary tokens" if targets.empty?

      path = CTC.forced_align(emissions[first_frame...last_frame, true], targets, blank: BLANK_ID)
      segments = CTC.merge_repeats(path, index_to_token)
      spans = CTC.spans(tokens_starred, segments, index_to_token.fetch(BLANK_ID))
      results = Text.postprocess(text_starred, spans, stride_ms)
      expected_tokens = PythonText.split(text)
      unless results.map { |word| word.fetch(:text) } == expected_tokens
        raise ArgumentError,
              "forced alignment did not preserve the complete ASR transcript " \
              "(#{results.length}/#{expected_tokens.length} words)"
      end
    rescue StandardError
      fallback_count += 1
      words.concat(uniform_fallback(text, start_time, end_time, segment_index))
      next
    end

    results.each_with_index do |word, word_index|
      absolute_start = (start_time + word.fetch(:start)).clamp(start_time, end_time)
      absolute_end = (start_time + word.fetch(:end)).clamp(absolute_start, end_time)
      words << TranscriptionWord.new(
        start: absolute_start,
        end: absolute_end,
        text: word.fetch(:text),
        segment_index: segment_index,
        segment_word_index: word_index,
        timing_source: "ctc"
      )
    end
  end
  [words.freeze, fallback_count]
end

#close ⇒ Object 



275
276
277
 
# File 'lib/cohere/transcribe/alignment/aligner.rb', line 275

def close
  session.close
end

#compute_emissions(audio) ⇒ Object 



289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
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
 
# File 'lib/cohere/transcribe/alignment/aligner.rb', line 289

def compute_emissions(audio)
  started = monotonic
  samples = mono_float32(audio)
  raise ArgumentError, "Cannot compute CTC emissions for empty audio" if samples.empty?
  return [compute_direct_emissions(samples), STRIDE_MS] if samples.length < WINDOW_SAMPLES

  total_windows = (samples.length + WINDOW_SAMPLES - 1) / WINDOW_SAMPLES
  extension_samples = (total_windows * WINDOW_SAMPLES) - samples.length
  extension_frames = extension_samples / INPUTS_TO_LOGITS_RATIO
  frame_count = (total_windows * WINDOW_FRAMES) - extension_frames
  raise TranscriptionRuntimeError, "MMS aligner produced no usable CTC frames" unless frame_count.positive?

  emissions = nil
  write_offset = 0
  first_window = 0
  current_batch_size = [batch_size, @learned_batch_size].min
  while first_window < total_windows
    window_count = [current_batch_size, total_windows - first_window].min
    begin
      input = build_window_batch(samples, first_window, window_count)
      logits = session.run(input)
      batch_log_probs = crop_and_normalize(logits, window_count)
    rescue NoMemoryError, StandardError => e
      raise unless out_of_memory?(e) && current_batch_size > 1

      current_batch_size = [1, current_batch_size / 2].max
      @learned_batch_size = current_batch_size
      next
    end

    expected_frames = window_count * WINDOW_FRAMES
    unless batch_log_probs.shape[0] == expected_frames
      raise TranscriptionRuntimeError, "MMS aligner returned an unexpected frame count"
    end

    class_count = batch_log_probs.shape[1]
    validate_class_count!(class_count)
    emissions ||= Numo::SFloat.zeros(frame_count, class_count + 1)
    if first_window + window_count == total_windows && extension_frames.positive?
      kept = batch_log_probs.shape[0] - extension_frames
      batch_log_probs = batch_log_probs[0...kept, true]
    end

    next_offset = write_offset + batch_log_probs.shape[0]
    raise TranscriptionRuntimeError, "MMS aligner produced too many CTC frames" if next_offset > emissions.shape[0]

    emissions[write_offset...next_offset, 0...class_count] = batch_log_probs
    write_offset = next_offset
    first_window += window_count
  end
  unless emissions && write_offset == emissions.shape[0]
    raise TranscriptionRuntimeError,
          "MMS emission assembly mismatch: wrote #{write_offset}, expected #{frame_count}"
  end

  [emissions, STRIDE_MS]
ensure
  @emissions_seconds += monotonic - started if started
end

#load! ⇒ Object 



270
271
272
273
 
# File 'lib/cohere/transcribe/alignment/aligner.rb', line 270

def load!
  session.load! if session.respond_to?(:load!)
  self
end

#load_seconds ⇒ Object 



266
267
268
 
# File 'lib/cohere/transcribe/alignment/aligner.rb', line 266

def load_seconds
  session.load_seconds
end

#provider ⇒ Object 



262
263
264
 
# File 'lib/cohere/transcribe/alignment/aligner.rb', line 262

def provider
  session.provider
end