cwenzi/neuroflow-cpp

1#include "neuroflow/model.hpp"2#include "neuroflow/generative.hpp"3#include "weight_io.hpp"4#include <iostream>5#include <fstream>6#include <sstream>7#include <random>8 9using namespace neuroflow;10 11NeuroFlowModel::Config load_config(const std::string& path) {12    NeuroFlowModel::Config cfg;13    std::ifstream f(path);14    if (!f) { std::cerr << "无法加载配置文件: " << path << std::endl; return cfg; }15    std::string json((std::istreambuf_iterator<char>(f)), std::istreambuf_iterator<char>());16 17    auto extract_num = [&](const std::string& key, size_t def = 0) {18        size_t p = json.find("\"" + key + "\"");19        if (p == std::string::npos) return def;20        p = json.find(':', p + key.size() + 2);21        while (p < json.size() && !std::isdigit(json[p])) p++;22        size_t e = p;23        while (e < json.size() && std::isdigit(json[e])) e++;24        return (e > p) ? std::stoul(json.substr(p, e - p)) : def;25    };26 27    cfg.vocab_size = extract_num("vocab_size", 5000);28    cfg.input_dim = extract_num("input_dim", 128);29    cfg.hidden_dim = extract_num("hidden_dim", 256);30    cfg.output_dim = extract_num("output_dim", cfg.vocab_size);31    cfg.num_layers = extract_num("num_layers", 2);32    cfg.memory_slots = extract_num("memory_slots", 64);33    cfg.memory_dim = extract_num("memory_dim", 128);34    cfg.num_associations = extract_num("num_associations", 8);35    cfg.use_causal_lm = true;36    cfg.max_seq_len = extract_num("max_seq_len", 128);37    cfg.causal_window_size = extract_num("causal_window_size", 32);38 39    std::cerr << "配置加载完成:" << std::endl;40    std::cerr << "  vocab=" << cfg.vocab_size << " d_model=" << cfg.input_dim41              << " hidden=" << cfg.hidden_dim << " output=" << cfg.output_dim << std::endl;42    return cfg;43}44 45// 判断是否为有效token（非特殊，且在词表范围内）46bool is_valid_token(size_t id, size_t vocab_actual) {47    return id >= 4 && id < vocab_actual;48}49 50int main(int argc, char* argv[]) {51    if (argc < 3) {52        std::cerr << "用法: " << argv[0] << " <config.json> <model.nfv1>" << std::endl;53        return 1;54    }55 56    std::string config_path = argv[1];57    std::string model_path = argv[2];58 59    // 哲学相关测试提示词60    std::vector<std::string> prompts = {61        "哲学",62        "辩证法",63        "唯物主义",64        "认识论",65        "存在",66        "意识",67        "真理",68        "实践",69    };70 71    std::cerr << "加载配置: " << config_path << std::endl;72    auto cfg = load_config(config_path);73 74    std::cerr << "构建模型..." << std::endl;75    NeuroFlowModel model(cfg);76 77    std::cerr << "加载权重: " << model_path << std::endl;78    model.load(model_path);79 80    auto stats = model.get_stats();81    std::cerr << "模型参数: " << stats.total_params << " 内存: " << stats.memory_bytes / 1024 / 1024 << " MB" << std::endl;82 83    std::string tok_path = config_path.substr(0, config_path.find_last_of("/\\") + 1) + "tokenizer_128k.json";84    std::cerr << "加载词表: " << tok_path << std::endl;85    BPETokenizer tokenizer(tok_path);86    size_t vocab_actual = tokenizer.vocab_size();87    std::cerr << "词表大小: " << vocab_actual << std::endl;88 89    float scale = 1.0f / cfg.vocab_size;90    std::mt19937 rng(42);91 92    for (auto& prompt : prompts) {93        std::cerr << "\n========================================\n";94        std::cerr << "提示词: " << prompt << std::endl;95        std::cerr << "========================================" << std::endl;96 97        std::vector<size_t> input_ids = tokenizer.encode(prompt);98        std::cerr << "输入tokens: ";99        for (auto id : input_ids) std::cerr << id << " ";100        std::cerr << std::endl;101 102        // 前向传播103        size_t seq_len = std::min(input_ids.size(), (size_t)cfg.max_seq_len);104        Tensor input({1, cfg.input_dim}, QuantType::FP32);105        float* inp = input.as_fp32();106        for (size_t j = 0; j < seq_len && j < cfg.input_dim; ++j) {107            inp[j] = static_cast<float>(input_ids[j]) * scale;108        }109 110        auto output = model.forward(input);111        const float* logits = output.output.as_fp32();112 113        // Top-10（只显示有效token）114        std::vector<std::pair<float, size_t>> scored;115        for (size_t i = 0; i < cfg.output_dim; ++i) {116            if (is_valid_token(i, vocab_actual))117                scored.push_back({logits[i], i});118        }119        std::sort(scored.begin(), scored.end(), std::greater<>());120 121        std::cout << "\n-- Top-10 有效token预测 --" << std::endl;122        for (int i = 0; i < std::min(10, (int)scored.size()); ++i) {123            size_t id = scored[i].second;124            float score = scored[i].first;125            std::string token = tokenizer.decode({id});126            std::cout << "  [" << id << "] \"" << token << "\" score=" << score << std::endl;127        }128 129        // 自回归生成（排除特殊token，使用温度采样）130        std::cout << "\n-- 自回归生成 --" << std::endl;131        std::vector<size_t> generated = input_ids;132        size_t last_id = input_ids.back();133        float temperature = 1.0f;134 135        for (int step = 0; step < 30; ++step) {136            Tensor step_input({1, cfg.input_dim}, QuantType::FP32);137            float* si = step_input.as_fp32();138            size_t ctx = std::min(generated.size(), (size_t)cfg.max_seq_len);139            size_t start = generated.size() - ctx;140            for (size_t j = 0; j < cfg.input_dim; ++j) {141                si[j] = (j < ctx) ? static_cast<float>(generated[start + j]) * scale : 0.0f;142            }143 144            auto out = model.forward(step_input);145            float* log = out.output.as_fp32();146 147            // 温度采样（只从有效token中选）148            float max_val = -1e30f;149            for (size_t j = 4; j < vocab_actual; ++j)150                if (log[j] > max_val) max_val = log[j];151 152            float sum_exp = 0.0f;153            std::vector<float> probs(cfg.output_dim, 0.0f);154            for (size_t j = 4; j < vocab_actual; ++j) {155                probs[j] = std::exp((log[j] - max_val) / temperature);156                sum_exp += probs[j];157            }158            for (size_t j = 4; j < vocab_actual; ++j)159                probs[j] /= sum_exp;160 161            // 累积采样162            float r = std::uniform_real_distribution<float>(0, 1)(rng);163            float cum = 0;164            size_t next_id = 4;165            for (size_t j = 4; j < vocab_actual; ++j) {166                cum += probs[j];167                if (r <= cum) { next_id = j; break; }168            }169 170            generated.push_back(next_id);171            std::string token = tokenizer.decode({next_id});172            std::cout << token;173 174            if (next_id == 2) break; // </s>175        }176        std::cout << std::endl;177    }178 179    return 0;180}181