main.cpp

/*'''
 * llm.cpp - LMGNU Organization
 * Copyright (C) 2026 Eamon Sippy
 * https://github.com/LMGNU/llm.cpp
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 */

#include "config/config.h"
#include "include/backward.h"
#include "include/llm-cpp.hpp"
#include "include/lm.h"
#include "include/sampler.h"
#include "include/tokenizer.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <csignal>
#include <cstdlib>
#include <ctime>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <string>
#include <vector>

#if defined(_WIN32)
#ifndef WIN32_LEAN_AND_MEAN
#define WIN32_LEAN_AND_MEAN
#endif
#include <psapi.h>
#include <windows.h>
#pragma comment(lib, "advapi32.lib")
#elif defined(__APPLE__)
#include <mach/mach.h>
#include <sys/sysctl.h>
#include <unistd.h>
#elif defined(__linux__)
#include <unistd.h>
#endif

static volatile bool g_interrupted = false;
static void sig_handler(int) { g_interrupted = true; }

static double wall_secs()
{
      using namespace std::chrono;
      return duration<double>(steady_clock::now().time_since_epoch()).count();
}

static std::string get_cpu_info()
{
#if defined(__linux__)
      std::ifstream cpuinfo("/proc/cpuinfo");
      std::string line;
      while (std::getline(cpuinfo, line))
      {
            if (line.find("model name") != std::string::npos)
            {
                  size_t colon = line.find(':');
                  if (colon != std::string::npos)
                  {
                        std::string name = line.substr(colon + 2);
                        if (name.length() > 34)
                              return name.substr(0, 31) + "...";
                        return name;
                  }
            }
      }
      return "Generic Linux CPU";
#elif defined(__APPLE__)
      char buffer[128];
      size_t bufferlen = sizeof(buffer);
      if (sysctlbyname("machdep.cpu.brand_string", &buffer, &bufferlen, NULL, 0) == 0)
      {
            std::string name(buffer);
            if (name.length() > 34)
                  return name.substr(0, 31) + "...";
            return name;
      }
      return "Apple Silicon / Mac CPU";
#elif defined(_WIN32)
      HKEY hKey;
      if (RegOpenKeyExA(HKEY_LOCAL_MACHINE, "HARDWARE\\DESCRIPTION\\System\\CentralProcessor\\0", 0,
                        KEY_READ, &hKey) == ERROR_SUCCESS)
      {
            char name[128];
            DWORD size = sizeof(name);
            if (RegQueryValueExA(hKey, "ProcessorNameString", NULL, NULL, (LPBYTE)name, &size) ==
                ERROR_SUCCESS)
            {
                  RegCloseKey(hKey);
                  std::string cpu_name(name);
                  if (cpu_name.length() > 34)
                        return cpu_name.substr(0, 31) + "...";
                  return cpu_name;
            }
            RegCloseKey(hKey);
      }
      return "Generic Windows CPU";
#else
      return "Generic CPU Host";
#endif
}

static double get_total_ram_mb()
{
#if defined(_WIN32)
      MEMORYSTATUSEX status;
      status.dwLength = sizeof(status);
      if (GlobalMemoryStatusEx(&status))
            return (double)status.ullTotalPhys / (1024.0 * 1024.0);
      return 0.0;
#elif defined(__APPLE__)
      int64_t mem = 0;
      size_t len = sizeof(mem);
      if (sysctlbyname("hw.memsize", &mem, &len, NULL, 0) == 0)
            return (double)mem / (1024.0 * 1024.0);
      return 0.0;
#elif defined(__linux__)
      std::ifstream meminfo("/proc/meminfo");
      std::string line;
      while (std::getline(meminfo, line))
      {
            if (line.compare(0, 9, "MemTotal:") == 0)
            {
                  std::istringstream iss(line.substr(9));
                  double kb = 0.0;
                  iss >> kb;
                  return kb / 1024.0;
            }
      }
      return 0.0;
#else
      return 0.0;
#endif
}

static double get_ram_usage_mb()
{
#if defined(_WIN32)
      PROCESS_MEMORY_COUNTERS pmc;
      if (K32GetProcessMemoryInfo(GetCurrentProcess(), &pmc, sizeof(pmc)))
            return (double)pmc.WorkingSetSize / (1024.0 * 1024.0);
      return 0.0;
#elif defined(__APPLE__)
      mach_task_basic_info_data_t info;
      mach_msg_type_number_t count = MACH_TASK_BASIC_INFO_COUNT;
      if (task_info(mach_task_self(), MACH_TASK_BASIC_INFO, (task_info_t)&info, &count) ==
          KERN_SUCCESS)
            return (double)info.resident_size / (1024.0 * 1024.0);
      return 0.0;
#elif defined(__linux__)
      std::ifstream status("/proc/self/status");
      std::string line;
      while (std::getline(status, line))
      {
            if (line.compare(0, 6, "VmRSS:") == 0)
            {
                  std::istringstream iss(line.substr(6));
                  double kb = 0.0;
                  iss >> kb;
                  return kb / 1024.0;
            }
      }
      return 0.0;
#else
      return 0.0;
#endif
}

static bool file_exists(const std::string &path)
{
      std::ifstream f(path.c_str(), std::ios::binary);
      return f.good();
}

static std::string dir_name(const std::string &path)
{
      std::string::size_type pos = path.find_last_of("/\\");
      if (pos == std::string::npos)
            return ".";
      if (pos == 0)
            return path.substr(0, 1);
      return path.substr(0, pos);
}

static bool is_absolute_path(const std::string &path)
{
      if (path.empty())
            return false;
      if (path.size() > 1 && path[1] == ':')
            return true;
      return path[0] == '/' || path[0] == '\\';
}

static std::string join_path(const std::string &base, const std::string &child)
{
      if (base.empty() || base == ".")
            return child;
      char last = base[base.size() - 1];
      if (last == '/' || last == '\\')
            return base + child;
      return base + "/" + child;
}

static std::string choose_existing_path(const std::string &requested_path, const std::string &argv0)
{
      if (requested_path.empty())
            return requested_path;
      if (file_exists(requested_path))
            return requested_path;
      if (is_absolute_path(requested_path))
            return requested_path;

      std::vector<std::string> candidates;
      candidates.push_back(join_path(dir_name(argv0), requested_path));
      candidates.push_back(join_path(".", requested_path));

      for (size_t i = 0; i < candidates.size(); ++i)
      {
            if (file_exists(candidates[i]))
                  return candidates[i];
      }
      return requested_path;
}

static std::string choose_output_path(const std::string &requested_path, const std::string &argv0)
{
      if (requested_path.empty() || is_absolute_path(requested_path))
            return requested_path;

      std::string exe_relative = join_path(dir_name(argv0), requested_path);
      if (file_exists(requested_path) || !file_exists(exe_relative))
            return requested_path;
      return exe_relative;
}

static void print_usage(const char *argv0)
{
      std::cout << "Usage: " << argv0 << " [options] [data_file]\n"
                << "\n"
                << "  --generate              run inference only (needs saved weights)\n"
                << "  --chat                  start interactive chat mode\n"
                << "  --chat-tokens N         max tokens per chat reply (default "
                << DEFAULT_CHAT_TOKENS << ")\n"
                << "  --system TEXT           system prompt prepended to every chat turn\n"
                << "  --rep-penalty F         repetition penalty, 1.0 means off (default "
                << DEFAULT_REP_PENALTY << ")\n"
                << "  --rep-window N          recent token window for penalty (default "
                << DEFAULT_REP_WINDOW << ")\n"
                << "  --help                  show this message\n";
}

static void sample_tokens(GPTLanguageModel &model, DataLoader &dl, int n_tokens,
                          const SamplerParams &params)
{
      std::vector<int> ctx = {0};
      for (int i = 0; i < n_tokens; ++i)
      {
            ctx = model.generate(ctx, 1, params);
            std::cout << dl.decode({ctx.back()}) << std::flush;
            if ((int)ctx.size() > BLOCK_SIZE)
                  ctx = std::vector<int>(ctx.end() - BLOCK_SIZE, ctx.end());
      }
      std::cout << "\n";
}

static float estimate_loss(GPTLanguageModel &model, DataLoader &dl, const std::string &split,
                           std::mt19937 &rng)
{
      float total = 0.0f;
      for (int k = 0; k < EVAL_ITERS; ++k)
      {
            std::pair<std::vector<int>, std::vector<int>> batch =
                dl.get_batch(split, BATCH_SIZE, BLOCK_SIZE, rng);
            std::pair<Tensor, float> result =
                model.forward(batch.first, BATCH_SIZE, BLOCK_SIZE, batch.second, false);
            total += result.second;
      }
      return total / EVAL_ITERS;
}

static float get_lr(int it, float max_lr, int max_iters)
{
      int warmup_iters = max_iters / 10;
      if (warmup_iters == 0)
            warmup_iters = 1;

      float min_lr = max_lr * 0.1f;

      if (it <= warmup_iters)
      {
            return max_lr * (float)it / (float)warmup_iters;
      }
      if (it > max_iters)
      {
            return min_lr;
      }

      float decay_ratio = (float)(it - warmup_iters) / (float)(max_iters - warmup_iters);
      float coeff = 0.5f * (1.0f + std::cos(3.14159265358979323846f * decay_ratio));

      return min_lr + coeff * (max_lr - min_lr);
}

static std::vector<int> build_turn_context(const std::vector<int> &sys_tokens,
                                           const std::vector<int> &user_tokens)
{
      std::vector<int> ctx;
      ctx.reserve(sys_tokens.size() + user_tokens.size());
      ctx.insert(ctx.end(), sys_tokens.begin(), sys_tokens.end());
      ctx.insert(ctx.end(), user_tokens.begin(), user_tokens.end());

      if ((int)ctx.size() > BLOCK_SIZE)
            ctx = std::vector<int>(ctx.end() - BLOCK_SIZE, ctx.end());

      return ctx;
}

static void run_chat(GPTLanguageModel &model, DataLoader &dl, int max_new_tokens,
                     const SamplerParams &params)
{
      std::vector<int> sys_tokens;
      if (!params.system_prompt.empty())
      {
            sys_tokens = dl.encode(params.system_prompt);
            if (sys_tokens.empty())
            {
                  std::cerr << "[WARN]  System prompt produced zero tokens.\n";
            }
            else
            {
                  std::cout << "[CHAT]  System prompt active (" << sys_tokens.size() << " tokens, "
                            << BLOCK_SIZE - (int)sys_tokens.size()
                            << " tokens left for user input)\n";
            }
      }

      std::cout << "\n" << std::string(60, '=') << "\n";
      std::cout << "  CHAT MODE\n";
      std::cout << "  Type your prompt and press Enter.\n";
      std::cout << "  Type quit or exit to leave.\n";
      std::cout << std::string(60, '=') << "\n\n";

      while (!g_interrupted)
      {
            std::cout << "root> ";
            std::cout.flush();

            std::string prompt;
            if (!std::getline(std::cin, prompt))
                  break;

            size_t s = prompt.find_first_not_of(" \t\r\n");
            size_t e = prompt.find_last_not_of(" \t\r\n");
            if (s == std::string::npos)
                  continue;
            prompt = prompt.substr(s, e - s + 1);

            if (prompt == "quit" || prompt == "exit")
            {
                  std::cout << "[Chat] Bye!\n";
                  break;
            }

            std::vector<int> user_tokens = dl.encode(prompt);
            if (user_tokens.empty())
                  user_tokens = {0};

            std::vector<int> ctx = build_turn_context(sys_tokens, user_tokens);

            std::cout << "llm> ";
            std::cout.flush();

            for (int tok = 0; tok < max_new_tokens && !g_interrupted; ++tok)
            {
                  ctx = model.generate(ctx, 1, params);
                  std::cout << dl.decode({ctx.back()}) << std::flush;

                  if ((int)ctx.size() > BLOCK_SIZE)
                        ctx = std::vector<int>(ctx.end() - BLOCK_SIZE, ctx.end());
            }
            std::cout << "\n\n";
      }
}

int main(int argc, char *argv[])
{
      std::signal(SIGINT, sig_handler);

      std::string data_path = DEFAULT_CLEANED_PATH;
      std::string model_path = BEST_MODEL_PATH;

      const char *env_data = std::getenv(DATA_PATH_ENV_VAR.c_str());
      const char *env_model = std::getenv(MODEL_PATH_ENV_VAR.c_str());
      if (env_data != nullptr && env_data[0] != '\0')
            data_path = env_data;
      if (env_model != nullptr && env_model[0] != '\0')
            model_path = env_model;

      bool gen_mode = false;
      bool chat_mode = false;
      int chat_tokens = DEFAULT_CHAT_TOKENS;
      float rep_penalty = DEFAULT_REP_PENALTY;
      int rep_window = DEFAULT_REP_WINDOW;
      std::string system_prompt;

      for (int i = 1; i < argc; ++i)
      {
            std::string a = argv[i];

            if (a == "--help")
            {
                  print_usage(argv[0]);
                  return 0;
            }
            else if (a == "--generate")
            {
                  gen_mode = true;
            }
            else if (a == "--chat")
            {
                  chat_mode = true;
            }
            else if (a == "--chat-tokens" && i + 1 < argc)
            {
                  chat_tokens = std::atoi(argv[++i]);
            }
            else if (a == "--system" && i + 1 < argc)
            {
                  system_prompt = argv[++i];
            }
            else if (a == "--rep-penalty" && i + 1 < argc)
            {
                  rep_penalty = (float)std::atof(argv[++i]);
            }
            else if (a == "--rep-window" && i + 1 < argc)
            {
                  rep_window = std::atoi(argv[++i]);
            }
            else
            {
                  data_path = a;
            }
      }

      data_path = choose_existing_path(data_path, argv[0]);
      model_path = choose_output_path(model_path, argv[0]);

      SamplerParams sampler;
      sampler.rep_penalty = rep_penalty;
      sampler.rep_window = rep_window;
      sampler.system_prompt = system_prompt;

      DataLoader dl;
      try
      {
            dl.load(data_path);
      }
      catch (const std::exception &e)
      {
            std::cerr << e.what() << "\n";
            std::cerr << "[HINT]  Put your text at " << DEFAULT_CLEANED_PATH
                      << ", pass a file path as the first argument, or set " << DATA_PATH_ENV_VAR
                      << ".\n";
            return 1;
      }

      GPTLanguageModel model(dl.vocab_size, N_EMBD, N_HEAD, N_LAYER, BLOCK_SIZE, SEED);

      long n_params = model.num_params();
      std::string cpu_spec = get_cpu_info();
      print_banner();
      std::string ram_spec = std::to_string((long long)get_total_ram_mb()) + " MB";

      std::cout << "\n";

      std::cout << "  "
                   "+------------------------------------------+-----------------------------------"
                   "-------+\n";
      std::cout << "  | " << std::left << std::setw(83) << "LLM Architecture" << " |\n";
      std::cout << "  "
                   "+------------------------------------------+-----------------------------------"
                   "-------+\n";
      std::cout << "  | Max Context Length   : " << std::left << std::setw(17) << BLOCK_SIZE
                << " | Vocab Size (BPE)     : " << std::left << std::setw(17) << dl.vocab_size
                << " |\n";
      std::cout << "  | Number of Layers     : " << std::left << std::setw(17) << N_LAYER
                << " | Attention Heads      : " << std::left << std::setw(17) << N_HEAD << " |\n";
      std::cout << "  | Embedding Channels   : " << std::left << std::setw(17) << N_EMBD
                << " | Total Parameters     : " << std::left << std::setw(17) << n_params << " |\n";
      std::cout << "  | Repetition Penalty   : " << std::left << std::setw(17) << rep_penalty
                << " | Repetition Window    : " << std::left << std::setw(17) << rep_window
                << " |\n";
      std::cout << "  "
                   "+------------------------------------------+-----------------------------------"
                   "-------+\n\n";

      std::cout << "  "
                   "+------------------------------------------------------------------------------"
                   "-------+\n";
      std::cout << "  | " << std::left << std::setw(83) << "Host Hardware Specs" << " |\n";
      std::cout << "  "
                   "+------------------------------------------------------------------------------"
                   "-------+\n";
      std::cout << "  | Host CPU Device      : " << std::left << std::setw(60) << cpu_spec
                << " |\n";
      std::cout << "  | Host RAM (Total)     : " << std::left << std::setw(60) << ram_spec
                << " |\n";
      std::cout << "  "
                   "+------------------------------------------------------------------------------"
                   "-------+\n\n";

      std::cout << std::right;

      if (chat_mode)
      {
            if (!file_exists(model_path))
            {
                  std::cerr << "[ERROR] Cannot start chat because model weights were not found at "
                            << model_path << "\n";
                  std::cerr << "[HINT]  Train first, or set " << MODEL_PATH_ENV_VAR
                            << " to an existing weights file.\n";
                  return 1;
            }

            model.load(model_path);
            std::cout << "weights:    " << model_path << "\n";
            std::cout << "max_tokens: " << chat_tokens << "\n";

            if (!sampler.system_prompt.empty())
                  std::cout << "system:     " << sampler.system_prompt << "\n";

            run_chat(model, dl, chat_tokens, sampler);
            return 0;
      }

      if (gen_mode)
      {
            if (!file_exists(model_path))
            {
                  std::cerr << "[ERROR] Cannot generate because model weights were not found at "
                            << model_path << "\n";
                  std::cerr << "[HINT]  Train first, or set " << MODEL_PATH_ENV_VAR
                            << " to an existing weights file.\n";
                  return 1;
            }

            model.load(model_path);
            std::cout << "\ngenerating:\n";
            std::vector<int> ctx = {0};
            while (!g_interrupted)
            {
                  ctx = model.generate(ctx, 1, sampler);
                  std::cout << dl.decode({ctx.back()}) << std::flush;
                  if ((int)ctx.size() > BLOCK_SIZE)
                        ctx = std::vector<int>(ctx.end() - BLOCK_SIZE, ctx.end());
            }
            std::cout << "\n";
            return 0;
      }

      // Train loop
      AdamWState opt = build_optimizer(model, LEARNING_RATE);
      std::mt19937 rng(SEED);

      float best_val_loss = 1e30f;
      float last_val_loss = 0.0f;

      {
            std::mt19937 init_rng(SEED);
            last_val_loss = estimate_loss(model, dl, "val", init_rng);
      }

      for (int iter = 1; iter <= MAX_ITERS && !g_interrupted; ++iter)
      {
            double step_start = wall_secs();

            float current_lr = get_lr(iter, LEARNING_RATE, MAX_ITERS);
            opt.lr = current_lr;

            std::pair<std::vector<int>, std::vector<int>> batch =
                dl.get_batch("train", BATCH_SIZE, BLOCK_SIZE, rng);

            SavedForward saved =
                forward_save(model, batch.first, BATCH_SIZE, BLOCK_SIZE, batch.second, true);
            float batch_loss =
                model.forward(batch.first, BATCH_SIZE, BLOCK_SIZE, batch.second, false).second;

            Grads grads = backward(model, saved);
            apply_grads(model, grads, opt);

            double step_ms = (wall_secs() - step_start) * 1000.0;
            int tok_per_sec =
                (step_ms > 0.0) ? (int)((long)BATCH_SIZE * BLOCK_SIZE / (step_ms / 1000.0)) : 0;

            bool val_updated = false;
            bool better = false;

            if (iter % EVAL_INTERVAL == 0 || iter == MAX_ITERS)
            {
                  last_val_loss = estimate_loss(model, dl, "val", rng);
                  val_updated = true;

                  if (last_val_loss < best_val_loss)
                  {
                        best_val_loss = last_val_loss;
                        model.save(model_path);
                        better = true;
                  }
            }

            double percent_done = ((double)iter / MAX_ITERS) * 100.0;
            double ram_mb = get_ram_usage_mb();

            std::cout << "step " << iter << "/" << MAX_ITERS << "(" << std::fixed
                      << std::setprecision(2) << percent_done << "%)"
                      << " | train loss " << std::fixed << std::setprecision(6) << batch_loss
                      << " | val loss " << std::fixed << std::setprecision(6) << last_val_loss
                      << (val_updated ? "*" : " ") << " | lr " << std::scientific
                      << std::setprecision(2) << current_lr << " | " << std::fixed
                      << std::setprecision(2) << std::setw(8) << step_ms << " ms"
                      << " | " << std::setw(6) << tok_per_sec << " tok/s"
                      << " | ram " << std::setprecision(1) << ram_mb << " MB"
                      << (better ? "  best" : "") << "\n";
            std::cout.flush();

            if (iter % EVAL_INTERVAL == 0 || iter == MAX_ITERS)
            {
                  std::cout << "generating:\n";
                  sample_tokens(model, dl, iter == MAX_ITERS ? 10000 : 150, sampler);
            }
      }

      return 0;
}