added loading models from assets

Author: Daniil Tcelikin

android/android.md

@@ -1,7 +1,10 @@
+import org.jetbrains.kotlin.de.undercouch.gradle.tasks.download.Download
+
 @Suppress("DSL_SCOPE_VIOLATION") // TODO: Remove once KTIJ-19369 is fixed
 plugins {
     alias(libs.plugins.androidApplication)
     alias(libs.plugins.kotlinAndroid)
+    alias(libs.plugins.download)
 }
 
 android {
@@ -29,9 +32,19 @@ android {
         sourceCompatibility = JavaVersion.VERSION_1_8
         targetCompatibility = JavaVersion.VERSION_1_8
     }
+
+    sourceSets {
+        maybeCreate("main").apply {
+            assets {
+                srcDirs("src/main/assets")
+            }
+        }
+    }
+
     kotlinOptions {
         jvmTarget = "1.8"
     }
+
     externalNativeBuild {
         cmake {
             path = file("src/main/cpp/CMakeLists.txt")
@@ -43,7 +56,32 @@ android {
     }
 }
 
+
+tasks {
+    val downloadTokenizer by creating(Download::class) {
+        onlyIf { !file("$projectDir/src/main/assets/tokenizer.bin").exists() }
+        src("https://github.com/karpathy/llama2.c/raw/master/tokenizer.bin")
+        dest("$projectDir/src/main/assets/tokenizer.bin")
+    }
+    val downloadModel by creating(Download::class) {
+        onlyIf { !file("$projectDir/src/main/assets/stories15M.bin").exists() }
+        src("https://huggingface.co/karpathy/tinyllamas/resolve/main/stories15M.bin")
+        dest("$projectDir/src/main/assets/stories15M.bin")
+    }
+    whenTaskAdded {
+        if (name in listOf("assembleDebug", "assembleRelease")) {
+            dependsOn(downloadTokenizer)
+            dependsOn(downloadModel)
+        }
+    }
+}
+
 dependencies {
+    implementation(libs.androidx.activity.ktx)
+    implementation(libs.androidx.lifecycle.viewmodel.ktx)
+    implementation(libs.androidx.lifecycle.livedata.ktx)
+    implementation(libs.androidx.lifecycle.runtime.ktx)
+
     implementation(libs.coroutines.core)
     implementation(libs.core.ktx)
     implementation(libs.appcompat)

android/app/build.gradle.kts

@@ -2,9 +2,11 @@ package com.celikin.llama2
 
 import androidx.appcompat.app.AppCompatActivity
 import android.os.Bundle
+import androidx.lifecycle.lifecycleScope
 import com.celikin.llama2.databinding.ActivityMainBinding
 import com.celikin.llama2.wrapper.InferenceRunner
 import com.celikin.llama2.wrapper.InferenceRunnerManager
+import kotlinx.coroutines.launch
 
 class MainActivity : AppCompatActivity() {
 
@@ -26,15 +28,21 @@ class MainActivity : AppCompatActivity() {
             val prompt = binding.promptEdit.text.toString()
             inferenceRunnerManager.run(prompt)
         }
-        initInference()
+
+        lifecycleScope.launch {
+            val assetsFolder = copyAssets(arrayOf("stories15M.bin", "tokenizer.bin"))
+            initInference(assetsFolder)
+        }
+
     }
 
     private fun updateText(token: String) {
         binding.sampleText.text = "${binding.sampleText.text}$token"
     }
 
-    private fun initInference() {
-        inferenceRunnerManager = InferenceRunnerManager().apply { init(callback) }
+    private fun initInference(assetsFolder: String) {
+        inferenceRunnerManager = InferenceRunnerManager()
+            .apply { init(callback, assetsFolder) }
     }
 
 }

android/app/src/main/java/com/celikin/llama2/MainActivity.kt

@@ -0,0 +1,36 @@
+package com.celikin.llama2
+
+import android.content.Context
+import android.util.Log
+import java.io.File
+import java.io.FileOutputStream
+
+
+private fun userAssetPath(context: Context?): String {
+    if (context == null)
+        return ""
+    val extDir = context.getExternalFilesDir("assets")
+        ?: return context.getDir("assets", 0).absolutePath
+    return extDir.absolutePath
+}
+
+fun Context.copyAssets(listFiles: Array<String>):String {
+    val extFolder = userAssetPath(this)
+    try {
+        assets.list("")
+            ?.filter { listFiles.contains(it) }
+            ?.filter { !File(extFolder, it).exists() }
+            ?.forEach {
+                val target = File(extFolder, it)
+                assets.open(it).use { input ->
+                    FileOutputStream(target).use { output ->
+                        input.copyTo(output)
+                        Log.i("Utils", "Copied from apk assets folder to ${target.absolutePath}")
+                    }
+                }
+            }
+    } catch (e: Exception) {
+        Log.e("Utils", "asset copy failed", e)
+    }
+    return extFolder
+}

android/app/src/main/java/com/celikin/llama2/Utils.kt

@@ -7,28 +7,30 @@ import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.launch
 
 class InferenceRunnerManager {
+    private lateinit var folderPath: String
     private val applicationScope = CoroutineScope(Dispatchers.IO + SupervisorJob())
 
-    fun init(callback: InferenceRunner.InferenceCallback) {
+    fun init(callback: InferenceRunner.InferenceCallback, folderPath: String) {
+        this.folderPath = folderPath
         InferenceRunner.setInferenceCallback(callback)
     }
 
     fun run(
         prompt: String = "",
         temperature: Float = 0.9f,
         steps: Int = 256,
-        checkpoint: String = "/data/local/tmp/stories15M.bin",
-        tokenizer: String = "/data/local/tmp/tokenizer.bin",
-        ompthreads: Int = 4,
+        checkpointFileName: String = "stories15M.bin",
+        tokenizerFileName: String = "tokenizer.bin",
+        ompThreads: Int = 4,
     ) {
         applicationScope.launch {
             InferenceRunner.run(
-                checkpoint = checkpoint,
-                tokenizer = tokenizer,
+                checkpoint = "$folderPath/$checkpointFileName",
+                tokenizer = "$folderPath/$tokenizerFileName",
                 temperature = temperature,
                 steps = steps,
                 prompt = prompt,
-                ompthreads = ompthreads
+                ompthreads = ompThreads
             )
         }
     }

android/app/src/main/java/com/celikin/llama2/wrapper/InferenceRunnerManager.kt

@@ -1,15 +1,22 @@
 [versions]
+activity-ktx = "1.7.2"
 agp = "8.2.0-alpha12"
 kotlin = "1.8.21"
-core-ktx = "1.9.0"
+core-ktx = "1.10.1"
+download = "5.3.1"
 junit = "4.13.2"
 androidx-test-ext-junit = "1.1.5"
 espresso-core = "3.5.1"
 appcompat = "1.6.1"
+lifecycle-viewmodel-ktx = "2.6.1"
 material = "1.9.0"
 constraintlayout = "2.1.4"
-kotlinx-coroutines-core = "1.6.4"
+kotlinx-coroutines-core = "1.7.1"
 [libraries]
+androidx-activity-ktx = { module = "androidx.activity:activity-ktx", version.ref = "activity-ktx" }
+androidx-lifecycle-livedata-ktx = { module = "androidx.lifecycle:lifecycle-livedata-ktx", version.ref = "lifecycle-viewmodel-ktx" }
+androidx-lifecycle-runtime-ktx = { module = "androidx.lifecycle:lifecycle-runtime-ktx", version.ref = "lifecycle-viewmodel-ktx" }
+androidx-lifecycle-viewmodel-ktx = { module = "androidx.lifecycle:lifecycle-viewmodel-ktx", version.ref = "lifecycle-viewmodel-ktx" }
 core-ktx = { group = "androidx.core", name = "core-ktx", version.ref = "core-ktx" }
 junit = { group = "junit", name = "junit", version.ref = "junit" }
 androidx-test-ext-junit = { group = "androidx.test.ext", name = "junit", version.ref = "androidx-test-ext-junit" }
@@ -22,4 +29,5 @@ coroutines-core = { group = "org.jetbrains.kotlinx", name = "kotlinx-coroutines-
 [plugins]
 androidApplication = { id = "com.android.application", version.ref = "agp" }
 kotlinAndroid = { id = "org.jetbrains.kotlin.android", version.ref = "kotlin" }
+download = { id = "de.undercouch.download", version.ref = "download" }