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)
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) 


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

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.



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

def batch_size
  @batch_size
end

#emissions_seconds ⇒ Object (readonly)

Returns the value of attribute emissions_seconds.



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

def emissions_seconds
  @emissions_seconds
end

#session ⇒ Object (readonly)

Returns the value of attribute session.



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

def session
  @session
end

#viterbi_seconds ⇒ Object (readonly)

Returns the value of attribute viterbi_seconds.



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

def viterbi_seconds
  @viterbi_seconds
end

Instance Method Details

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



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

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:



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

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 



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

def close
  session.close
end

#compute_emissions(audio) ⇒ Object 



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

def compute_emissions(audio)
  started = monotonic
  samples = mono_float32(audio)
  raise ArgumentError, "Cannot compute CTC emissions for empty audio" if samples.empty?

  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]
    unless class_count == VOCABULARY.length
      raise TranscriptionRuntimeError,
            "MMS aligner returned #{class_count} classes; expected #{VOCABULARY.length}"
    end
    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 



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

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

#load_seconds ⇒ Object 



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

def load_seconds
  session.load_seconds
end

#provider ⇒ Object 



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

def provider
  session.provider
end