#include <iostream>
#include <string>
#include <justchat/chat.hpp>



int main(int argc, char **argv) {
    // Check usage
    if (argc != 2) {
        std::cerr << "Usage: " << argv[0] << " <model config path>" << std::endl;
        return EXIT_FAILURE;
    }
    // Get args
    const auto model_config = argv[1];
    // Construct chat model
    LM::Chat::Inference model(model_config);
    if (!model.start()) {
        std::cerr << "Failed to load model." << std::endl;
        return EXIT_FAILURE;
    }
    std::string instruction;
    for (;;) {
        std::cout << "> ";
        std::getline(std::cin, instruction);
        model.prompt(instruction, [] (auto token) {
            std::cout << token << std::flush;
            return true;
        }, [] (float progress, bool) {
            std::cout << ' ' << unsigned(progress) << "%  \r" << std::flush;
            return true;
        });
        std::cout << '\n';
    }
}