lmtps

#!/usr/bin/env -S java --enable-native-access=ALL-UNNAMED --source 25 -cp ./zbo/lightmetal.jar:./lightmetal.jar

import lm.catalog.boundary.ModelCatalog;
import lm.catalog.control.ModelsDirectory;
import lm.configuration.control.ZCfg;
import lm.configuration.entity.GenerationConfig;
import lm.generation.boundary.LightMetal;
import lm.generation.entity.Tps;

String name = MethodHandles.lookup().lookupClass().getName();
String prompt = "Count from 1 to 10.";

void main(String... args) throws Exception {
    System.setProperty("progress", "false");
    ZCfg.load("lightmetal");

    var all = Stream.of(args).anyMatch("-all"::equals);

    if (all) {
        var models = ModelCatalog.list();
        if (models.isEmpty()) {
            System.err.println(name + ": no models found in " + ModelsDirectory.path());
            System.exit(1);
        }
        models.forEach(m -> measure(m, prompt));
        return;
    }

    var fileName = ZCfg.string("model");
    if (fileName == null) {
        System.err.println(name + ": no model configured (set model=<file> in ~/.lightmetal/app.properties)");
        System.exit(1);
    }
    if (!measure(fileName, prompt)) {
        System.exit(1);
    }
}

boolean measure(String fileName, String prompt) {
    var modelPath = ModelCatalog.resolve(fileName);
    if (!Files.exists(modelPath)) {
        System.err.println(name + ": model not found at " + modelPath);
        return false;
    }
    var count = new long[1];
    var firstTokenNanos = new long[1];
    try (var lm = LightMetal.load(modelPath)) {
        var modelName = lm.metadata().name().orElse(modelPath.getFileName().toString());
        try (var stream = lm.generate(prompt, GenerationConfig.defaults())) {
            stream.forEach(t -> {
                if (firstTokenNanos[0] == 0L) firstTokenNanos[0] = System.nanoTime();
                count[0]++;
                IO.print(t.text());
                System.out.flush();
            });
        }
        IO.println("");
        if (count[0] < 2) {
            System.err.println(name + ": " + modelName + ": got " + count[0] + " tokens, need at least 2 for tps");
            return false;
        }
        IO.println("[" + modelName + ": " + Tps.measure(count[0], firstTokenNanos[0]) + "]");
    }
    return true;
}