#include <cstdio>
#include <print>
#include <string>
#include <vector>

#include "mie/chat_templates/Message.h"
#include "mie/chat_templates/Role.h"
#include "mie/llms/Llama3.h"
#include "mie/models/KVCache.h"
#include "mie/samplers/RandomSampler.h"

int main() {
    mie::Llama3<float> llm;

    if (!mie::load("path/to/Llama-3.2-1B-Instruct-FP32", llm)) {
        // error: failed to load llm
    }

    mie::RandomSampler<float> sampler;
    sampler.config().temperature = 0.6;
    sampler.config().topP = 0.95;
    sampler.config().topK = 50;

    mie::KVCache<float> cache;

    std::vector<mie::Message> messages{
        {mie::Role::kSystem, "You are a helpful assistant."},
        {mie::Role::kUser, "How do I make a cake?"},
    };

    for (const std::string &token : llm(messages, cache, sampler)) {
        std::print("{}", token);
        fflush(stdout);
    }
}

#include <cstdint>
#include <cstdio>
#include <print>
#include <string>
#include <vector>

#include <Eigen/Dense>

#include "mie/chat_templates/Llama3ChatTemplate.h"
#include "mie/chat_templates/Message.h"
#include "mie/chat_templates/Role.h"
#include "mie/models/KVCache.h"
#include "mie/models/Llama3Model.h"
#include "mie/samplers/RandomSampler.h"
#include "mie/tokenizers/Llama3Tokenizer.h"

using Vector = Eigen::RowVector<float, Eigen::Dynamic>;

int main() {
    mie::Llama3Model<float> model;

    if (!mie::load("path/to/Llama-3.2-1B-Instruct-FP32", model)) {
        // error: failed to load model
    }

    mie::Llama3Tokenizer tokenizer;

    if (!mie::load("path/to/Llama-3.2-1B-Instruct-FP32", tokenizer)) {
        // error: failed to load tokenizer
    }

    mie::Llama3ChatTemplate chatTemplate;

    mie::RandomSampler<float> sampler;
    sampler.config().temperature = 0.6;
    sampler.config().topP = 0.95;
    sampler.config().topK = 50;

    mie::KVCache<float> cache;

    std::vector<mie::Message> messages{
        {mie::Role::kSystem, "You are a helpful assistant."},
        {mie::Role::kUser, "How do I make a cake?"},
    };

    std::string input = chatTemplate(messages);

    std::vector<uint32_t> inputIds = tokenizer(input);

    for (;;) {
        Vector logits = model(inputIds, cache).bottomRows(1);

        uint32_t id = sampler(logits);

        std::string token = tokenizer.convertIdToToken(id);

        std::print("{}", token);
        fflush(stdout);

        inputIds.push_back(id);
    }
}