-
Notifications
You must be signed in to change notification settings - Fork 75
/
GPT2Client.kt
216 lines (182 loc) · 7.96 KB
/
GPT2Client.kt
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
package co.huggingface.android_transformers.gpt2.ml
import android.app.Application
import android.text.Spannable
import android.text.SpannableStringBuilder
import android.util.JsonReader
import android.widget.TextView
import androidx.core.content.res.ResourcesCompat
import androidx.databinding.BindingAdapter
import androidx.lifecycle.*
import co.huggingface.android_transformers.gpt2.R
import co.huggingface.android_transformers.gpt2.tokenization.GPT2Tokenizer
import kotlinx.coroutines.*
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.random.Random
private const val SEQUENCE_LENGTH = 64
private const val VOCAB_SIZE = 50257
private const val NUM_HEAD = 12
private const val NUM_LITE_THREADS = 4
private const val MODEL_PATH = "model.tflite"
private const val VOCAB_PATH = "gpt2-vocab.json"
private const val MERGES_PATH = "gpt2-merges.txt"
private const val TAG = "GPT2Client"
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 val initJob: Job
private var autocompleteJob: Job? = null
private lateinit var tokenizer: GPT2Tokenizer
private lateinit var tflite: Interpreter
private val prompts = arrayOf(
"Before boarding your rocket to Mars, remember to pack these items",
"In a shocking finding, scientist discovered a herd of unicorns living in a remote, previously unexplored valley, in the Andes Mountains. Even more surprising to the researchers was the fact that the unicorns spoke perfect English.",
"Legolas and Gimli advanced on the orcs, raising their weapons with a harrowing war cry.",
"Today, scientists confirmed the worst possible outcome: the massive asteroid will collide with Earth",
"Hugging Face is a company that releases awesome projects in machine learning because"
)
private val _prompt = MutableLiveData(prompts.random())
val prompt: LiveData<String> = _prompt
private val _completion = MutableLiveData("")
val completion: LiveData<String> = _completion
var strategy = GPT2Strategy(GPT2StrategyEnum.TOPK, 40)
init {
initJob = viewModelScope.launch {
val encoder = loadEncoder()
val decoder = encoder.entries.associateBy({ it.value }, { it.key })
val bpeRanks = loadBpeRanks()
tokenizer = GPT2Tokenizer(encoder, decoder, bpeRanks)
tflite = loadModel()
}
}
override fun onCleared() {
super.onCleared()
tflite.close()
}
fun launchAutocomplete() {
autocompleteJob = viewModelScope.launch {
initJob.join()
autocompleteJob?.cancelAndJoin()
_completion.value = ""
generate(_prompt.value!!)
}
}
fun refreshPrompt() {
_prompt.value = prompts.random()
launchAutocomplete()
}
private suspend fun generate(text: String, nbTokens: Int = 100) = withContext(Dispatchers.Default) {
val tokens = tokenizer.encode(text)
repeat (nbTokens) {
val maxTokens = tokens.takeLast(SEQUENCE_LENGTH).toIntArray()
val paddedTokens = maxTokens + IntArray(SEQUENCE_LENGTH - maxTokens.size)
val inputIds = Array(1) { paddedTokens }
val predictions: Predictions = Array(1) { Array(SEQUENCE_LENGTH) { FloatArray(VOCAB_SIZE) } }
val outputs = mutableMapOf<Int, Any>(0 to predictions)
tflite.runForMultipleInputsOutputs(arrayOf(inputIds), outputs)
val outputLogits = predictions[0][maxTokens.size-1]
val nextToken: Int = when (strategy.strategy) {
GPT2StrategyEnum.TOPK -> {
val filteredLogitsWithIndexes = outputLogits
.mapIndexed { index, fl -> (index to fl) }
.sortedByDescending { it.second }
.take(strategy.value)
// Softmax computation on filtered logits
val filteredLogits = filteredLogitsWithIndexes.map { it.second }
val maxLogitValue = filteredLogits.max()!!
val logitsExp = filteredLogits.map { exp(it - maxLogitValue) }
val sumExp = logitsExp.sum()
val probs = logitsExp.map { it.div(sumExp) }
val logitsIndexes = filteredLogitsWithIndexes.map { it.first }
sample(logitsIndexes, probs)
}
else -> outputLogits.argmax()
}
tokens.add(nextToken)
val decodedToken = tokenizer.decode(listOf(nextToken))
_completion.postValue(_completion.value + decodedToken)
yield()
}
}
private suspend fun loadModel(): Interpreter = withContext(Dispatchers.IO) {
val assetFileDescriptor = getApplication<Application>().assets.openFd(MODEL_PATH)
assetFileDescriptor.use {
val fileChannel = FileInputStream(assetFileDescriptor.fileDescriptor).channel
val modelBuffer = fileChannel.map(FileChannel.MapMode.READ_ONLY, it.startOffset, it.declaredLength)
val opts = Interpreter.Options()
opts.setNumThreads(NUM_LITE_THREADS)
return@use Interpreter(modelBuffer, opts)
}
}
private suspend fun loadEncoder(): Map<String, Int> = withContext(Dispatchers.IO) {
hashMapOf<String, Int>().apply {
val vocabStream = getApplication<Application>().assets.open(VOCAB_PATH)
vocabStream.use {
val vocabReader = JsonReader(InputStreamReader(it, "UTF-8"))
vocabReader.beginObject()
while (vocabReader.hasNext()) {
val key = vocabReader.nextName()
val value = vocabReader.nextInt()
put(key, value)
}
vocabReader.close()
}
}
}
private suspend fun loadBpeRanks():Map<Pair<String, String>, Int> = withContext(Dispatchers.IO) {
hashMapOf<Pair<String, String>, Int>().apply {
val mergesStream = getApplication<Application>().assets.open(MERGES_PATH)
mergesStream.use { stream ->
val mergesReader = BufferedReader(InputStreamReader(stream))
mergesReader.useLines { seq ->
seq.drop(1).forEachIndexed { i, s ->
val list = s.split(" ")
val keyTuple = list[0] to list[1]
put(keyTuple, i)
}
}
}
}
}
}
private fun randomIndex(probs: List<Float>): Int {
val rnd = probs.sum() * 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) {
if (this[it] > this[bestIndex]) {
bestIndex = it
}
}
return bestIndex
}
@BindingAdapter("prompt", "completion")
fun TextView.formatCompletion(prompt: String, completion: String): Unit {
text = when {
completion.isEmpty() -> prompt
else -> {
val str = SpannableStringBuilder(prompt + completion)
val bgCompletionColor = ResourcesCompat.getColor(resources, R.color.colorPrimary, context.theme)
str.setSpan(android.text.style.BackgroundColorSpan(bgCompletionColor), prompt.length, str.length, Spannable.SPAN_EXCLUSIVE_EXCLUSIVE)
str
}
}
}