app/src/main/kotlin/com/honjimaku/subrep/Whisper.kt (4273 bytes)
1 package com.honjimaku.subrep 2 3 import java.io.File 4 5 /** JNI surface. See src/main/cpp/whisper_jni.c. */ 6 internal object WhisperLib { 7 external fun initContext(modelPath: String): Long 8 external fun freeContext(ptr: Long) 9 external fun transcribe(ptr: Long, samples: FloatArray, threads: Int, language: String, audioCtx: Int): Int 10 external fun segmentCount(ptr: Long): Int 11 external fun segmentText(ptr: Long, index: Int): ByteArray 12 external fun detectedLanguage(ptr: Long): String 13 external fun systemInfo(): String 14 } 15 16 class WhisperUnavailable(message: String) : Exception(message) 17 18 /** A loaded speech model. Not thread-safe: one transcription at a time. */ 19 class Whisper private constructor(private var ptr: Long) : Recognizer { 20 21 /** The language that Whisper found in the last piece: "ja", "ru", ... */ 22 override val detectedLanguage: String get() = WhisperLib.detectedLanguage(ptr) 23 24 /** 25 * The text of one piece of sound: [samples] at 16 kHz mono, from -1 to 1. [language] is a 26 * code such as "ja", or "auto". 27 */ 28 override fun transcribe(samples: FloatArray, language: String): String { 29 val rc = WhisperLib.transcribe(ptr, samples, threads, language, audioContext(samples.size)) 30 if (rc != 0) throw WhisperUnavailable("The speech model failed on this sound (code $rc).") 31 return (0 until WhisperLib.segmentCount(ptr)).joinToString(" ") { 32 // Lenient decode: a small model can end a segment inside a character. 33 String(WhisperLib.segmentText(ptr, it), Charsets.UTF_8).replace("�", "").trim() 34 }.trim() 35 } 36 37 override fun close() { 38 if (ptr != 0L) WhisperLib.freeContext(ptr) 39 ptr = 0 40 } 41 42 companion object { 43 44 @Volatile private var loaded = false 45 46 val threads: Int by lazy { fastCores() } 47 48 fun open(model: File): Whisper { 49 requireCpu() 50 if (!loaded) { 51 System.loadLibrary("subrep_whisper") 52 loaded = true 53 } 54 val ptr = WhisperLib.initContext(model.absolutePath) 55 if (ptr == 0L) throw WhisperUnavailable("The speech model could not be loaded.") 56 return Whisper(ptr) 57 } 58 59 /** 60 * The part of the encoder window to compute for a piece of [samples]: 50 units for each 61 * second of sound, plus a margin. The whole window is 1500 units for 30 s. A piece of 62 * 5 s then costs a fifth of the encoder, which is most of the time on a phone. 63 */ 64 fun audioContext(samples: Int): Int = (samples / 16_000.0 * 50 + 100).toInt().coerceIn(384, 1500) 65 66 /** 67 * The native library is built for ARMv8.2 with half-precision and dot-product 68 * instructions. Loading it without them is a crash with no message, so look first. 69 */ 70 private fun requireCpu() { 71 val features = runCatching { 72 File("/proc/cpuinfo").readLines() 73 .firstOrNull { it.startsWith("Features") }.orEmpty() 74 .substringAfter(':').trim().split(' ').toSet() 75 }.getOrDefault(emptySet()) 76 if (features.isEmpty()) return 77 val missing = listOf("asimdhp", "asimddp").filter { it !in features } 78 if (missing.isNotEmpty()) { 79 throw WhisperUnavailable("The processor of this phone is too old for speech recognition on the phone (ARMv8.2 is needed).") 80 } 81 } 82 83 /** 84 * Threads to use: the fast cores only. Whisper waits for its slowest thread at each 85 * layer, so a little core makes it slower, not faster. "Fast" is a core clocked above 86 * the slowest cluster. 87 */ 88 private fun fastCores(): Int { 89 val max = (0 until Runtime.getRuntime().availableProcessors()).mapNotNull { cpu -> 90 runCatching { 91 File("/sys/devices/system/cpu/cpu$cpu/cpufreq/cpuinfo_max_freq").readText().trim().toLong() 92 }.getOrNull() 93 } 94 if (max.isEmpty()) return 4 95 val slowest = max.min() 96 val fast = max.count { it > slowest } 97 return if (fast == 0) (max.size / 2).coerceIn(2, 4) else fast.coerceIn(2, 6) 98 } 99 } 100 }