|
| 1 | +package sk.ainet.apps.skainet.cli |
| 2 | + |
| 3 | +import sk.ainet.apps.kllama.CpuAttentionBackend |
| 4 | +import sk.ainet.apps.kllama.cli.AgentCli |
| 5 | +import sk.ainet.apps.kllama.cli.ToolCallingDemo |
| 6 | +import sk.ainet.apps.llm.InferenceRuntime |
| 7 | +import sk.ainet.apps.llm.Tokenizer |
| 8 | +import sk.ainet.apps.llm.UnifiedModelLoader |
| 9 | +import sk.ainet.apps.llm.generate |
| 10 | +import sk.ainet.apps.llm.backend.BackendRegistry |
| 11 | +import sk.ainet.apps.llm.backend.bestAvailable |
| 12 | +import sk.ainet.apps.llm.tokenizer.TokenizerFactory |
| 13 | +import sk.ainet.apps.kllama.chat.ModelMetadata |
| 14 | +import sk.ainet.context.DirectCpuExecutionContext |
| 15 | +import sk.ainet.io.JvmRandomAccessSource |
| 16 | +import sk.ainet.io.model.QuantPolicy |
| 17 | +import sk.ainet.lang.tensor.data.MemorySegmentTensorDataFactory |
| 18 | +import sk.ainet.lang.types.FP32 |
| 19 | +import sk.ainet.models.llama.LlamaRuntime |
| 20 | +import sk.ainet.models.llama.LlamaWeightLoader |
| 21 | +import sk.ainet.models.llama.LlamaWeightMapper |
| 22 | +import sk.ainet.models.llama.MemSegWeightConverter |
| 23 | +import java.lang.foreign.Arena |
| 24 | +import java.nio.file.Path |
| 25 | +import kotlin.io.path.exists |
| 26 | +import kotlin.io.path.extension |
| 27 | +import kotlin.system.exitProcess |
| 28 | +import kotlinx.coroutines.runBlocking |
| 29 | +import kotlin.time.measureTime |
| 30 | + |
| 31 | +private data class CliArgs( |
| 32 | + val modelPath: Path, |
| 33 | + val steps: Int, |
| 34 | + val temperature: Float, |
| 35 | + val prompt: String?, |
| 36 | + val chatMode: Boolean, |
| 37 | + val agentMode: Boolean, |
| 38 | + val demoMode: Boolean, |
| 39 | + val templateName: String?, |
| 40 | + val contextLength: Int? |
| 41 | +) |
| 42 | + |
| 43 | +private fun usage(errorMessage: String? = null): Nothing { |
| 44 | + if (errorMessage != null) { |
| 45 | + System.err.println("Error: $errorMessage") |
| 46 | + System.err.println() |
| 47 | + } |
| 48 | + |
| 49 | + println("Usage: skainet -m <model.gguf> [-s <steps>] [-k <temperature>] [--chat] [--agent] [--demo] [--template=NAME] <prompt>") |
| 50 | + println() |
| 51 | + println(" -m, --model Path to .gguf model (required)") |
| 52 | + println(" -s, --steps Generation steps (default: 64)") |
| 53 | + println(" -k, --temperature Sampling temperature (default: 0.8)") |
| 54 | + println(" --chat Interactive chat mode") |
| 55 | + println(" --agent Interactive agent mode with tool calling") |
| 56 | + println(" --demo Tool calling demo with file listing and calculator") |
| 57 | + println(" --template=NAME Chat template: llama3, chatml, qwen, gemma (auto-detected if omitted)") |
| 58 | + println(" --context=N Cap context length to N tokens") |
| 59 | + println(" -h, --help Show this help") |
| 60 | + println() |
| 61 | + println("Supported architectures (auto-detected from GGUF metadata):") |
| 62 | + println(" LLaMA, Mistral, Qwen2, Qwen3, Gemma, Apertus") |
| 63 | + println() |
| 64 | + println("Examples:") |
| 65 | + println(" skainet -m model.gguf \"The capital of France is\"") |
| 66 | + println(" skainet -m model.gguf --chat") |
| 67 | + println(" skainet -m model.gguf --demo \"What is 2 + 2?\"") |
| 68 | + exitProcess(if (errorMessage == null) 0 else 1) |
| 69 | +} |
| 70 | + |
| 71 | +private fun parseArgs(args: Array<String>): CliArgs { |
| 72 | + if (args.isEmpty()) usage("Missing arguments.") |
| 73 | + |
| 74 | + var model: String? = null |
| 75 | + var steps = 64 |
| 76 | + var temperature = 0.8f |
| 77 | + var prompt: String? = null |
| 78 | + var chatMode = false |
| 79 | + var agentMode = false |
| 80 | + var demoMode = false |
| 81 | + var templateName: String? = null |
| 82 | + var contextLength: Int? = null |
| 83 | + |
| 84 | + var idx = 0 |
| 85 | + fun nextValue(flag: String): String { |
| 86 | + if (idx + 1 >= args.size) usage("$flag requires a value.") |
| 87 | + return args[++idx] |
| 88 | + } |
| 89 | + |
| 90 | + while (idx < args.size) { |
| 91 | + val arg = args[idx] |
| 92 | + when { |
| 93 | + arg == "-h" || arg == "--help" -> usage() |
| 94 | + arg == "-m" || arg == "--model" -> model = nextValue(arg) |
| 95 | + arg.startsWith("--model=") -> model = arg.substringAfter("=") |
| 96 | + arg == "-s" || arg == "--steps" -> { |
| 97 | + val value = nextValue(arg) |
| 98 | + steps = value.toIntOrNull() ?: usage("Invalid steps '$value'.") |
| 99 | + } |
| 100 | + arg == "-k" || arg == "--temperature" -> { |
| 101 | + val value = nextValue(arg) |
| 102 | + temperature = value.toFloatOrNull() ?: usage("Invalid temperature '$value'.") |
| 103 | + } |
| 104 | + arg == "--chat" -> chatMode = true |
| 105 | + arg == "--agent" -> agentMode = true |
| 106 | + arg == "--demo" -> demoMode = true |
| 107 | + arg.startsWith("--template=") -> templateName = arg.substringAfter("=") |
| 108 | + arg.startsWith("--context=") -> { |
| 109 | + val value = arg.substringAfter("=") |
| 110 | + contextLength = value.toIntOrNull() ?: usage("Invalid context length '$value'.") |
| 111 | + } |
| 112 | + arg.startsWith("-") -> usage("Unknown option '$arg'.") |
| 113 | + else -> { |
| 114 | + if (prompt != null) usage("Multiple prompts provided.") |
| 115 | + prompt = arg |
| 116 | + } |
| 117 | + } |
| 118 | + idx++ |
| 119 | + } |
| 120 | + |
| 121 | + val modelPath = model?.let { Path.of(it) } ?: usage("Model is required (-m/--model).") |
| 122 | + |
| 123 | + if (!chatMode && !agentMode && !demoMode && prompt == null) { |
| 124 | + usage("Prompt is required (or use --chat/--agent/--demo mode).") |
| 125 | + } |
| 126 | + |
| 127 | + return CliArgs(modelPath, steps, temperature, prompt, chatMode, agentMode, demoMode, templateName, contextLength) |
| 128 | +} |
| 129 | + |
| 130 | +fun main(args: Array<String>) { |
| 131 | + runBlocking { |
| 132 | + val cliArgs = parseArgs(args) |
| 133 | + val modelPath = cliArgs.modelPath |
| 134 | + |
| 135 | + if (!modelPath.exists()) error("Model not found: $modelPath") |
| 136 | + if (modelPath.extension.lowercase() != "gguf") { |
| 137 | + error("Only GGUF models are supported by the unified CLI. Use model-specific CLIs for other formats.") |
| 138 | + } |
| 139 | + |
| 140 | + // Auto-detect architecture |
| 141 | + val modelInfo = UnifiedModelLoader.peek { JvmRandomAccessSource.open(modelPath.toString()) } |
| 142 | + println("Architecture: ${modelInfo.architecture}, Family: ${modelInfo.family.displayName}") |
| 143 | + println("Dimensions: ${modelInfo.embeddingLength}d, ${modelInfo.blockCount} layers, vocab=${modelInfo.vocabSize}") |
| 144 | + |
| 145 | + // Select backend |
| 146 | + val provider = BackendRegistry.bestAvailable() |
| 147 | + println("Backend: ${provider.displayName}") |
| 148 | + |
| 149 | + // Set up execution context |
| 150 | + val quantArena = Arena.ofShared() |
| 151 | + val memSegFactory = MemorySegmentTensorDataFactory() |
| 152 | + val ctx = DirectCpuExecutionContext(tensorDataFactory = memSegFactory) |
| 153 | + |
| 154 | + Runtime.getRuntime().addShutdownHook(Thread { |
| 155 | + quantArena.close() |
| 156 | + memSegFactory.close() |
| 157 | + }) |
| 158 | + |
| 159 | + // Load model based on detected family |
| 160 | + val acceptedArchitectures = modelInfo.family.architectures + setOf(modelInfo.architecture) |
| 161 | + val loader = LlamaWeightLoader( |
| 162 | + randomAccessProvider = { JvmRandomAccessSource.open(modelPath.toString()) }, |
| 163 | + quantPolicy = QuantPolicy.NATIVE_OPTIMIZED, |
| 164 | + acceptedArchitectures = acceptedArchitectures |
| 165 | + ) |
| 166 | + |
| 167 | + println("Loading GGUF model from $modelPath (${modelInfo.family.displayName}, streaming)...") |
| 168 | + val loaded = loader.loadToMapStreaming<FP32, Float>(ctx, FP32::class) |
| 169 | + val rawWeights = LlamaWeightMapper.map(loaded) |
| 170 | + |
| 171 | + val runtimeWeights = if (rawWeights.quantTypes.isNotEmpty()) { |
| 172 | + println("Converting ${rawWeights.quantTypes.size} quantized tensors to SIMD format...") |
| 173 | + MemSegWeightConverter.convert(rawWeights, ctx, quantArena) |
| 174 | + } else { |
| 175 | + rawWeights |
| 176 | + } |
| 177 | + |
| 178 | + if (cliArgs.contextLength != null) { |
| 179 | + println("Context length capped to ${cliArgs.contextLength} (model default: ${runtimeWeights.metadata.contextLength})") |
| 180 | + } |
| 181 | + |
| 182 | + val backend = CpuAttentionBackend<FP32>( |
| 183 | + ctx, runtimeWeights, FP32::class, |
| 184 | + ropeFreqBase = runtimeWeights.metadata.ropeFreqBase, |
| 185 | + maxContextLength = cliArgs.contextLength |
| 186 | + ) |
| 187 | + |
| 188 | + @Suppress("DEPRECATION") |
| 189 | + val runtime: InferenceRuntime<FP32> = LlamaRuntime<FP32>( |
| 190 | + ctx, runtimeWeights, backend, FP32::class, |
| 191 | + eps = runtimeWeights.metadata.rmsNormEps |
| 192 | + ) |
| 193 | + |
| 194 | + // Load tokenizer from GGUF |
| 195 | + println("Loading embedded GGUF tokenizer...") |
| 196 | + val tokenizer: Tokenizer = JvmRandomAccessSource.open(modelPath.toString()).use { source -> |
| 197 | + TokenizerFactory.fromGGUF(source) |
| 198 | + } |
| 199 | + |
| 200 | + // Build model metadata for chat template auto-detection |
| 201 | + val metadata = ModelMetadata( |
| 202 | + family = modelInfo.family.id, |
| 203 | + architecture = modelInfo.architecture, |
| 204 | + chatTemplate = modelInfo.fields["tokenizer.chat_template"] as? String |
| 205 | + ) |
| 206 | + |
| 207 | + // Dispatch |
| 208 | + if (cliArgs.chatMode || cliArgs.agentMode || cliArgs.demoMode) { |
| 209 | + when { |
| 210 | + cliArgs.demoMode -> { |
| 211 | + val demo = ToolCallingDemo(runtime, tokenizer, cliArgs.templateName, metadata) |
| 212 | + demo.run(maxTokens = cliArgs.steps, temperature = cliArgs.temperature) |
| 213 | + } |
| 214 | + cliArgs.agentMode -> { |
| 215 | + val agentCli = AgentCli(runtime, tokenizer, cliArgs.templateName, metadata) |
| 216 | + agentCli.runAgent(maxTokens = cliArgs.steps, temperature = cliArgs.temperature) |
| 217 | + } |
| 218 | + else -> { |
| 219 | + val agentCli = AgentCli(runtime, tokenizer, cliArgs.templateName, metadata) |
| 220 | + agentCli.runChat(maxTokens = cliArgs.steps, temperature = cliArgs.temperature) |
| 221 | + } |
| 222 | + } |
| 223 | + return@runBlocking |
| 224 | + } |
| 225 | + |
| 226 | + // Standard generation mode |
| 227 | + val promptText = cliArgs.prompt ?: error("Prompt is required for standard generation mode.") |
| 228 | + val promptTokens = tokenizer.encode(promptText) |
| 229 | + |
| 230 | + println("Generating ${cliArgs.steps} tokens with temperature=${cliArgs.temperature}...") |
| 231 | + println("---") |
| 232 | + print(promptText) |
| 233 | + |
| 234 | + val elapsed = measureTime { |
| 235 | + runtime.generate(prompt = promptTokens, steps = cliArgs.steps, temperature = cliArgs.temperature) { id -> |
| 236 | + print(tokenizer.decode(id)) |
| 237 | + } |
| 238 | + }.inWholeMilliseconds |
| 239 | + |
| 240 | + val tokPerSec = cliArgs.steps / elapsed.toDouble() * 1000 |
| 241 | + println("\n---") |
| 242 | + println("tok/s: $tokPerSec") |
| 243 | + } |
| 244 | +} |
0 commit comments