gpt2 generation with topk

Author: Pierrci

app/src/main/java/co/huggingface/android_transformers/gpt2/ml/GPT2Client.kt

@@ -3,15 +3,15 @@ package co.huggingface.android_transformers.gpt2.ml
 import android.app.Application
 import android.util.JsonReader
 import androidx.lifecycle.AndroidViewModel
-import androidx.lifecycle.liveData
-import androidx.lifecycle.viewModelScope
 import co.huggingface.android_transformers.gpt2.tokenization.GPT2Tokenizer
-import kotlinx.coroutines.Dispatchers
 import org.tensorflow.lite.Interpreter
 import java.io.BufferedReader
 import java.io.FileInputStream
 import java.io.InputStreamReader
 import java.nio.channels.FileChannel
+import kotlin.math.exp
+import kotlin.math.min
+import kotlin.random.Random
 
 private const val SEQUENCE_LENGTH  = 64
 private const val VOCAB_SIZE       = 50257
@@ -23,10 +23,15 @@ private const val MERGES_PATH      = "gpt2-merges.txt"
 
 private typealias Predictions = Array<Array<FloatArray>>
 
+enum class GPT2StrategyEnum { GREEDY, TOPK }
+data class GPT2Strategy(val strategy: GPT2StrategyEnum, val value: Int = 0)
+
 class GPT2Client(application: Application) : AndroidViewModel(application) {
     private lateinit var tokenizer: GPT2Tokenizer
     private lateinit var tflite: Interpreter
 
+    var strategy = GPT2Strategy(GPT2StrategyEnum.TOPK, 40)
+
     fun init() {
         if (!::tokenizer.isInitialized) {
             val encoder  = loadEncoder()
@@ -46,7 +51,7 @@ class GPT2Client(application: Application) : AndroidViewModel(application) {
     fun generate(text: String, nbTokens: Int = 10) { // = liveData<String>(
             //viewModelScope.coroutineContext+Dispatchers.Default) {
 
-        var tokens = tokenizer.encode(text)
+        val tokens = tokenizer.encode(text)
         repeat (nbTokens) {
             val maxTokens    = tokens.takeLast(SEQUENCE_LENGTH).toIntArray()
             val paddedTokens = maxTokens + IntArray(SEQUENCE_LENGTH - maxTokens.size)
@@ -57,7 +62,26 @@ class GPT2Client(application: Application) : AndroidViewModel(application) {
 
             tflite.runForMultipleInputsOutputs(arrayOf(inputIds), outputs)
             val outputLogits = predictions[0][maxTokens.size-1]
-            val nextToken    = outputLogits.argmax()
+
+            val nextToken: Int = when (strategy.strategy) {
+                GPT2StrategyEnum.TOPK -> {
+                    val finalTopK = min(strategy.value, outputLogits.size)
+                    val filteredLogits = outputLogits
+                            .mapIndexed { index, fl -> (index to fl) }
+                            .sortedBy   { it.second }
+                            .takeWhile  { it.second < finalTopK }
+
+                    // Softmax computation on filtered logits
+                    val maxLogitValue = outputLogits.max()!!
+                    val logitsExp     = filteredLogits.map { exp(it.second - maxLogitValue) }
+                    val sumExp        = logitsExp.sum()
+                    val probs         = logitsExp.map { it.div(sumExp) }
+
+                    val logitsIndexes = filteredLogits.map { it.first }
+                    sample(logitsIndexes, probs)
+                }
+                else -> outputLogits.argmax()
+            }
 
             tokens.add(nextToken)
             val decodedToken = tokenizer.decode(listOf(nextToken))
@@ -111,6 +135,25 @@ class GPT2Client(application: Application) : AndroidViewModel(application) {
     }
 }
 
+private fun randomIndex(probs: List<Float>): Int {
+    val rnd = Random.nextFloat()
+    var acc = 0f
+
+    probs.forEachIndexed { i, fl ->
+        acc += fl
+        if (rnd < acc) {
+            return i
+        }
+    }
+
+    return probs.size - 1
+}
+
+private fun sample(indexes: List<Int>, probs: List<Float>): Int {
+    val i = randomIndex(probs)
+    return indexes[i]
+}
+
 private fun FloatArray.argmax(): Int {
     var bestIndex = 0
     repeat(size) {

app/src/main/java/co/huggingface/android_transformers/gpt2/tokenization/GPT2Tokenizer.kt

@@ -13,8 +13,8 @@ class GPT2Tokenizer(
     }
 
     fun encode(text: String): MutableList<Int> {
-        val tokens = encodeRegex.findAll(text).map {
-            it.value.codePoints()
+        val tokens = encodeRegex.findAll(text).map { result ->
+            result.value.codePoints()
                     .boxed()
                     .map { byteEncoder[it]!! }
                     .toArray()