cohere-transcribe 0.1.0

data/lib/cohere/transcribe/alignment/ctc.rb

@@ -0,0 +1,224 @@
+# frozen_string_literal: true
+
+require "numo/narray"
+
+module Cohere
+  module Transcribe
+    module Alignment
+      Segment = Data.define(:label, :start, :end)
+
+      # Ruby port of TorchAudio's CPU forced_align kernel. The two-row score
+      # matrix remains float32 so its Viterbi choices follow the reference
+      # kernel's rounding and tie-breaking behavior, while the back-pointers
+      # use one byte per state and frame.
+      module CTC
+        module_function
+
+        def forced_align(log_probs, targets, blank:)
+          frame_count, class_count = matrix_shape(log_probs)
+          validate_targets!(targets, blank, class_count)
+
+          target_count = targets.length
+          repeat_count = targets.each_cons(2).count { |left, right| left == right }
+          minimum_frames = target_count + repeat_count
+          if frame_count < minimum_frames
+            raise ArgumentError,
+                  "targets length is too long for CTC: " \
+                  "#{frame_count} frames, #{target_count} targets, #{repeat_count} repeats"
+          end
+
+          state_count = (2 * target_count) + 1
+          alphas = Numo::SFloat.zeros(2, state_count).fill(-Float::INFINITY)
+          back_pointers = "\xFF".b * (frame_count * state_count)
+
+          first_state = frame_count > minimum_frames ? 0 : 1
+          end_state = state_count == 1 ? 1 : 2
+          (first_state...end_state).each do |state|
+            alphas[0, state] = probability(log_probs, 0, label_for(state, targets, blank))
+          end
+
+          1.upto(frame_count - 1) do |frame|
+            if frame_count - frame <= minimum_frames
+              first_state += 1 if first_state.odd? && targets.fetch(first_state / 2) != targets.fetch((first_state / 2) + 1)
+              first_state += 1
+            end
+            if frame <= minimum_frames
+              if end_state.even? && end_state < 2 * target_count &&
+                 targets.fetch((end_state / 2) - 1) != targets.fetch(end_state / 2)
+                end_state += 1
+              end
+              end_state += 1
+            end
+
+            current_row = frame % 2
+            previous_row = (frame - 1) % 2
+            alphas[current_row, true] = -Float::INFINITY
+            loop_start = first_state
+            if first_state.zero?
+              alphas[current_row, 0] = alphas[previous_row, 0] + probability(log_probs, frame, blank)
+              back_pointers.setbyte(frame * state_count, 0)
+              loop_start += 1
+            end
+
+            (loop_start...end_state).each do |state|
+              stay = alphas[previous_row, state]
+              advance = alphas[previous_row, state - 1]
+              skip = -Float::INFINITY
+              if state.odd? && state != 1 && targets.fetch(state / 2) != targets.fetch((state / 2) - 1)
+                skip = alphas[previous_row, state - 2]
+              end
+
+              score, step = if skip > advance && skip > stay
+                              [skip, 2]
+                            elsif advance > stay && advance > skip
+                              [advance, 1]
+                            else
+                              [stay, 0]
+                            end
+              back_pointers.setbyte((frame * state_count) + state, step)
+              alphas[current_row, state] =
+                score + probability(log_probs, frame, label_for(state, targets, blank))
+            end
+          end
+
+          final_row = (frame_count - 1) % 2
+          state = if alphas[final_row, state_count - 1] > alphas[final_row, state_count - 2]
+                    state_count - 1
+                  else
+                    state_count - 2
+                  end
+          path = Array.new(frame_count)
+          (frame_count - 1).downto(0) do |frame|
+            path[frame] = label_for(state, targets, blank)
+            break if frame.zero?
+
+            step = back_pointers.getbyte((frame * state_count) + state)
+            raise "CTC back-pointer was not initialized" unless step && step <= 2
+
+            state -= step
+          end
+          path.freeze
+        end
+
+        def merge_repeats(path, index_to_token)
+          segments = []
+          first = 0
+          while first < path.length
+            following = first + 1
+            following += 1 while following < path.length && path.fetch(first) == path.fetch(following)
+            label = index_to_token.fetch(path.fetch(first))
+            segments << Segment.new(label: label, start: first, end: following - 1)
+            first = following
+          end
+          segments.freeze
+        end
+
+        def spans(tokens, segments, blank)
+          letter_index = 0
+          token_index = 0
+          intervals = []
+          interval_start = 0
+
+          segments.each_with_index do |segment, segment_index|
+            if token_index == tokens.length
+              unless segment_index == segments.length - 1 && segment.label == blank
+                raise "CTC path continued with a non-blank label after the transcript"
+              end
+
+              next
+            end
+
+            token_letters = tokens.fetch(token_index).split
+            letter = token_letters.fetch(letter_index)
+            next if segment.label == blank
+            raise "#{segment.label} != #{letter}" unless segment.label == letter
+
+            interval_start = segment_index if letter_index.zero?
+            if letter_index == token_letters.length - 1
+              letter_index = 0
+              token_index += 1
+              intervals << [interval_start, segment_index]
+              while token_index < tokens.length && tokens.fetch(token_index).empty?
+                intervals << [segment_index, segment_index]
+                token_index += 1
+              end
+            else
+              letter_index += 1
+            end
+          end
+          raise "CTC path did not consume every transcript token" unless token_index == tokens.length
+
+          intervals.each_with_index.map do |(first, last), interval_index|
+            span = segments[first..last]
+            if first.positive?
+              previous = segments.fetch(first - 1)
+              if previous.label == blank
+                padded_start = if interval_index.zero?
+                                 previous.start
+                               else
+                                 ((previous.start + previous.end) / 2.0).to_i
+                               end
+                span = [Segment.new(label: blank, start: padded_start, end: span.first.start), *span]
+              end
+            end
+            if last + 1 < segments.length
+              following = segments.fetch(last + 1)
+              if following.label == blank
+                padded_end = if interval_index == intervals.length - 1
+                               following.end
+                             else
+                               ((following.start + following.end) / 2.0).floor
+                             end
+                span = [*span, Segment.new(label: blank, start: span.last.end, end: padded_end)]
+              end
+            end
+            span.freeze
+          end.freeze
+        end
+
+        def matrix_shape(values)
+          if values.respond_to?(:ndim)
+            raise ArgumentError, "CTC emissions must be a two-dimensional matrix" unless values.ndim == 2
+
+            frames, classes = values.shape
+          elsif values.is_a?(Array) && values.all?(Array)
+            frames = values.length
+            classes = values.first&.length || 0
+            raise ArgumentError, "CTC emissions rows must have equal widths" unless values.all? { |row| row.length == classes }
+          else
+            raise ArgumentError, "CTC emissions must be an indexable matrix"
+          end
+          raise ArgumentError, "CTC emissions must contain at least one frame" unless frames.positive?
+          raise ArgumentError, "CTC emissions must contain at least two classes" unless classes >= 2
+
+          [frames, classes]
+        end
+        private_class_method :matrix_shape
+
+        def validate_targets!(targets, blank, class_count)
+          unless blank.is_a?(Integer) && blank.between?(0, class_count - 1)
+            raise ArgumentError, "blank must be within the emissions vocabulary"
+          end
+          unless targets.is_a?(Array) && !targets.empty? && targets.all?(Integer)
+            raise ArgumentError, "CTC targets must be a non-empty integer array"
+          end
+          raise ArgumentError, "CTC targets must not contain blank index #{blank}" if targets.include?(blank)
+          return if targets.all? { |target| target.between?(0, class_count - 1) }
+
+          raise ArgumentError, "CTC targets must be within the emissions vocabulary"
+        end
+        private_class_method :validate_targets!
+
+        def label_for(state, targets, blank)
+          state.even? ? blank : targets.fetch(state / 2)
+        end
+        private_class_method :label_for
+
+        def probability(values, frame, label)
+          values.respond_to?(:ndim) ? values[frame, label] : values.fetch(frame).fetch(label)
+        end
+        private_class_method :probability
+      end
+    end
+  end
+end

data/lib/cohere/transcribe/alignment/text.rb

@@ -0,0 +1,237 @@
+# frozen_string_literal: true
+
+require_relative "../python_text"
+require_relative "uroman_data"
+
+module Cohere
+  module Transcribe
+    module Alignment
+      # Arabic/English normalization used by the pinned MMS aligner. The
+      # punctuation inventory mirrors the retained Python reference, including
+      # its corpus-derived private-use and invalid-codepoint entries.
+      module Text
+        mms_punctuation_hex = <<~HEX
+          0008 001A 0020 0023-0026 0028-002B 002F 0033 0035 0037 003D
+          0040-0043 0045-0049 004B-004D 0051-0055 0057 005B 005D-005F
+          007C 007E 0081 008D 008F-0090 0092-0094 009D 00A3 00A9 00AC
+          00AF-00B0 00B6-00B7 00BC-00BD 00D7 02C7 02DC 037E 0387 05BE
+          061E 06D4 06D6-06D7 06DA 0965 09F7 0E46 0EC6 0F04-0F05
+          0F0C-0F0E 104A-104F 109F 1361-1366 166D-166E 17D4-17D7 17DA
+          2006 2015 2020-2022 2026 202F-2030 2038 203B 2044 2060 206F
+          2071 207F 20AC 2153 215F 2192 2212 226A-226B 230A-230B
+          231E-231F 2D70 3005 3010-3011 3014-3015 30FB-30FC A4FE-A4FF
+          A92E-A92F E003 E013 E01D E040-E041 E043 E06B E06D E070 E203
+          E231-E234 E236 E290-E296 E2D9 E2DB E2DE E2F9-E2FC E2FE E313
+          E328-E32E E492-E493 E514 E516-E518 EA01 EA7B F171-F173 F218
+          F21D F50E-F50F F511 F513 F518-F519 F521-F523 F529 F52D F62D
+          F7D0 F7E8 FD3E-FD3F FE50 FE55-FE57 FF0E FF1C FF1E FF65 FFFD
+        HEX
+        MMS_PUNCTUATION_HEX = mms_punctuation_hex.tr("\n", " ").freeze
+
+        PREFIX_CODEPOINTS = [
+          *".?,:!{}“”«»‹›„‚‟‛’‘><¿。;՚՛՜՝՞՟։¡،、।\"؛؟，！？；：（）﹏⋯、‧／～─＿".codepoints,
+          *(0x300C..0x300F), *(0xFF41..0xFF44), *(0x3008..0x300B)
+        ].freeze
+
+        PUNCTUATION_CODEPOINTS = begin
+          values = MMS_PUNCTUATION_HEX.split.flat_map do |token|
+            first, last = token.split("-", 2).map { |hex| Integer(hex, 16) }
+            last ? (first..last).to_a : [first]
+          end
+          (values + PREFIX_CODEPOINTS).to_h { |codepoint| [codepoint, true] }.freeze
+        end
+
+        DELETION_CODEPOINTS = [
+          0x200E, 0x200C, *(0x0656..0x0657), 0x200B, *(0x064B..0x0652),
+          0x202C, 0x200F, 0x202A
+        ].to_h { |codepoint| [codepoint, true] }.freeze
+
+        ARABIC_ROMANIZATION = {
+          "ء" => "'", "آ" => "a", "أ" => "a", "ؤ" => "w", "إ" => "i", "ئ" => "ye",
+          "ا" => "a", "ب" => "b", "ة" => "a", "ت" => "t", "ث" => "th", "ج" => "j",
+          "ح" => "h", "خ" => "kh", "د" => "d", "ذ" => "th", "ر" => "r", "ز" => "z",
+          "س" => "s", "ش" => "sh", "ص" => "s", "ض" => "d", "ط" => "t", "ظ" => "z",
+          "ع" => "'", "غ" => "gh", "ف" => "f", "ق" => "q", "ك" => "k", "ل" => "l",
+          "م" => "m", "ن" => "n", "ه" => "h", "و" => "w", "ى" => "a", "ي" => "y",
+          "ٮ" => "b", "ٯ" => "q", "ٹ" => "tt", "ٺ" => "tt", "ٻ" => "b", "ټ" => "t",
+          "ٽ" => "t", "پ" => "p", "ٿ" => "t", "ڀ" => "b", "ځ" => "z", "ڂ" => "h",
+          "ڃ" => "ny", "ڄ" => "dy", "څ" => "ts", "چ" => "tch", "ڇ" => "tch",
+          "ڈ" => "dd", "ډ" => "d", "ڊ" => "d", "ڋ" => "d", "ڌ" => "d", "ڍ" => "dd",
+          "ڎ" => "d", "ڏ" => "d", "ڐ" => "d", "ڑ" => "rr", "ژ" => "j", "ښ" => "kh",
+          "ڤ" => "v", "ک" => "k", "ڪ" => "k", "ګ" => "g", "ڭ" => "ng", "گ" => "g",
+          "ں" => "n", "ھ" => "h", "ہ" => "h", "ۃ" => "a", "ۇ" => "u", "ۈ" => "yu",
+          "ۋ" => "v", "ی" => "i", "ۍ" => "y", "ێ" => "y", "ې" => "e", "ے" => "y",
+          "ە" => "ae", "ٴ" => "h"
+        }.freeze
+
+        LATIN_ROMANIZATION = {
+          "æ" => "ae", "œ" => "oe", "ø" => "oe", "ß" => "ss", "ð" => "d", "þ" => "th",
+          "ł" => "l", "đ" => "d", "ħ" => "h", "ı" => "i", "ə" => "e"
+        }.freeze
+
+        # Multi-codepoint rules from Uroman 1.3.1.1 that are reachable from the
+        # pinned ASR tokenizers. Uroman chooses these before its single-character
+        # candidates; keeping them here preserves that longest-match behavior.
+        UROMAN_SEQUENCE_ROMANIZATION = {
+          "λμπ" => "lb", "νμπ" => "nb", "ρμπ" => "rb",
+          "γγ" => "ng", "ει" => "ei", "ευ" => "eu", "αυ" => "au", "ου" => "ou",
+          "ηυ" => "eu", "υι" => "ui", "ωυ" => "ou", "ντ" => "nd",
+          "シャ" => "sha", "シュ" => "shu", "ショ" => "sho",
+          "チャ" => "cha", "チェ" => "che", "チュ" => "chu", "チョ" => "cho",
+          "ジャ" => "ja", "ジュ" => "ju", "ジョ" => "jo", "ジェ" => "je",
+          "ヂャ" => "ja", "ヂュ" => "ju", "ヂョ" => "jo",
+          "フェ" => "fe", "ヴェ" => "ve", "フィ" => "fi", "ウィ" => "wi",
+          "ヴィ" => "vi", "ティ" => "ti", "ディ" => "di",
+          "しゃ" => "sha", "しゅ" => "shu", "しょ" => "sho",
+          "ちゃ" => "cha", "ちゅ" => "chu", "ちょ" => "cho",
+          "じゃ" => "ja", "じゅ" => "ju", "じょ" => "jo",
+          "ぢゃ" => "ja", "ぢゅ" => "ju", "ぢょ" => "jo"
+        }.freeze
+        UROMAN_SEQUENCE_LENGTHS = UROMAN_SEQUENCE_ROMANIZATION.keys.map(&:length).uniq.sort.reverse.freeze
+        JAPANESE_SMALL_Y = "ゃゅょャュョ".chars.freeze
+        JAPANESE_SMALL_TSU = "っッ".chars.freeze
+        JAPANESE_SCRIPT = {
+          "ゃ" => :hiragana, "ゅ" => :hiragana, "ょ" => :hiragana,
+          "ャ" => :katakana, "ュ" => :katakana, "ョ" => :katakana
+        }.freeze
+
+        DIGIT_PATTERN = "0-9০-৯០-៩०-९୦-୯۰-۹꤀-꤉０-９൦-൯၀-၉ⅰ-ⅹ"
+
+        module_function
+
+        def normalize(text)
+          normalized = PythonText.nfkc_lower(text.to_s)
+          normalized = normalized.gsub(/\([^)]*\p{Nd}[^)]*\)/, " ")
+          normalized = normalized.gsub("&lt;", "").gsub("&gt;", "").gsub("&nbsp", "")
+          normalized = normalized.gsub(/(\S+)[\u201B\u2019\u2018](\S+)/, "\\1'\\2")
+          normalized = normalized.tr("ٱٰۥۦ", "ااوي")
+          normalized = normalized.delete("ـٓ")
+          normalized = normalized.tr("ٕٔ", "ءء")
+          normalized = normalized.each_char.map do |character|
+            PUNCTUATION_CODEPOINTS.key?(character.ord) ? " " : character
+          end.join
+          normalized = normalized.each_char.reject do |character|
+            DELETION_CODEPOINTS.key?(character.ord)
+          end.join
+          whitespace = PythonText.whitespace_class
+          normalized = normalized.gsub(/\A[#{DIGIT_PATTERN}]+(?=#{whitespace})/, " ")
+          normalized = normalized.gsub(/(?<=#{whitespace})[#{DIGIT_PATTERN}]+(?=#{whitespace}|\z)/, " ")
+          PythonText.collapse(normalized)
+        end
+
+        def romanize(text, language)
+          normalized = normalize(text)
+          transliterated = romanize_normalized(normalized, language)
+          spaced = PythonText.strip(transliterated).each_char.to_a.join(" ").downcase
+          PythonText.strip(spaced.gsub(/[^a-z' ]/, " ").gsub(/ +/, " "))
+        end
+
+        def preprocess(text, language)
+          words = PythonText.split(text.to_s)
+          tokens = words.map { |word| romanize(word, language) }
+          token_stream = []
+          text_stream = []
+          tokens.zip(words).each do |token, word|
+            token_stream.push("<star>", token)
+            text_stream.push("<star>", word)
+          end
+          [token_stream.freeze, text_stream.freeze]
+        end
+
+        def postprocess(text_starred, spans, stride_ms)
+          results = []
+          text_starred.each_with_index do |text, index|
+            next if text == "<star>"
+
+            span = spans.fetch(index)
+            results << {
+              start: span.first.start * stride_ms / 1_000.0,
+              end: (span.last.end + 1) * stride_ms / 1_000.0,
+              text: text
+            }
+          end
+          results.each_cons(2) do |current, following|
+            following[:start] = current[:end] if following[:start] < current[:end]
+          end
+          results.freeze
+        end
+
+        def romanize_character(character, _language)
+          mapped = UROMAN_SINGLE_CODEPOINT[character.ord]
+          return mapped unless mapped.nil?
+
+          return character if /[a-z']/i.match?(character)
+          return ARABIC_ROMANIZATION.fetch(character) if ARABIC_ROMANIZATION.key?(character)
+          return LATIN_ROMANIZATION.fetch(character) if LATIN_ROMANIZATION.key?(character)
+
+          decomposed = character.unicode_normalize(:nfkd).gsub(/\p{Mn}/, "")
+          decomposed.match?(/\A[a-z]+\z/i) ? decomposed : ""
+        end
+        private_class_method :romanize_character
+
+        def romanize_normalized(normalized, language)
+          characters = normalized.each_char.to_a
+          output = +""
+          index = 0
+          while index < characters.length
+            if JAPANESE_SMALL_TSU.include?(characters[index]) && index + 1 < characters.length
+              following, = romanize_unit(characters, index + 1, language)
+              consonant = following.match(/\A(ch|[bcdfghjklmnpqrstwz])/)&.[](1)
+              if consonant
+                output << (consonant == "ch" ? "t" : consonant)
+                index += 1
+                next
+              end
+            end
+
+            romanized, consumed = romanize_unit(characters, index, language)
+            output << romanized
+            index += consumed
+          end
+          output
+        end
+        private_class_method :romanize_normalized
+
+        def romanize_unit(characters, index, language)
+          UROMAN_SEQUENCE_LENGTHS.each do |length|
+            sequence = characters[index, length]&.join
+            mapped = UROMAN_SEQUENCE_ROMANIZATION[sequence]
+            return [mapped, length] unless mapped.nil?
+          end
+
+          pair = characters[index, 2]&.join
+          if pair == "μπ"
+            return [word_start?(characters, index) ? "b" : "mb", 2]
+          elsif pair == "γκ"
+            return [word_start?(characters, index) ? "g" : "ng", 2]
+          end
+
+          character = characters.fetch(index)
+          mapped = romanize_character(character, language)
+          following = characters[index + 1]
+          if JAPANESE_SMALL_Y.include?(following) && same_japanese_script?(character, following) &&
+             mapped.match?(/[bcdfghjklmnpqrstvwxyz]i\z/)
+            return [mapped.delete_suffix("i") + romanize_character(following, language), 2]
+          end
+
+          [mapped, 1]
+        end
+        private_class_method :romanize_unit
+
+        def word_start?(characters, index)
+          index.zero? || !characters[index - 1].match?(/\p{L}/)
+        end
+        private_class_method :word_start?
+
+        def same_japanese_script?(character, following)
+          script = JAPANESE_SCRIPT[following]
+          return false if script.nil?
+
+          codepoint = character.ord
+          script == :hiragana ? codepoint.between?(0x3040, 0x309F) : codepoint.between?(0x30A0, 0x30FF)
+        end
+        private_class_method :same_japanese_script?
+      end
+    end
+  end
+end