user_entry.cpp

//
// Created by wangchen on 2024/5/3.
//
#include <iostream>
#include <string>
#include <chrono>
#include "src/utils/model_utils.h"

using namespace llm;

void PrintRes(int index, const char* content) {
    if (index == -1) {
        std::cout << std::endl;
    } else {
        std::cout << content << std::flush;
    }
}

int main() {
    std::cout << "=== LLaMA 多轮对话系统 (mmap权重加载) ===" << std::endl;
    
    std::string model_dir = "/llm/nankai/wanghui_space/llama2/weight_llama2/";
    std::string tokenizer_file = "/llm/nankai/wanghui_space/llama2/weight_llama2/tokenizer.bin";
    
    std::cout << "正在创建模型..." << std::endl;
    auto start_time = std::chrono::high_resolution_clock::now();
    
    // 创建模型（自动使用mmap模式）
    auto model = CreateRealLLMModel<float>(model_dir, tokenizer_file);
    
    auto end_time = std::chrono::high_resolution_clock::now();
    auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time);
    
    std::cout << "模型创建完成，耗时: " << duration.count() << " 毫秒" << std::endl;
    std::cout << "\n=== 开始多轮对话 ===" << std::endl;
    std::cout << "输入 'quit' 或 'exit' 退出对话" << std::endl;
    std::cout << "输入 'clear' 清空对话历史" << std::endl;
    std::cout << "----------------------------------------" << std::endl;
    
    std::string history = "";
    int round = 0;
    
    while (true) {
        std::cout << "\n用户: ";
        std::string input;
        std::getline(std::cin, input);
        
        // 检查退出命令
        if (input == "quit" || input == "exit") {
            std::cout << "再见！" << std::endl;
            break;
        }
        
        // 检查清空历史命令
        if (input == "clear") {
            history = "";
            round = 0;
            std::cout << "对话历史已清空" << std::endl;
            continue;
        }
        
        // 跳过空输入
        if (input.empty()) {
            continue;
        }
        
        // 生成输入数据
        auto input_data = model->MakeInput(history, round, input);
        
        std::cout << "助手: ";
        auto inference_start = std::chrono::high_resolution_clock::now();
        
        // 进行推理
        std::string response = model->Response(input_data, PrintRes);
        
        auto inference_end = std::chrono::high_resolution_clock::now();
        auto inference_duration = std::chrono::duration_cast<std::chrono::milliseconds>(inference_end - inference_start);
        
        std::cout << "\n[推理耗时: " << inference_duration.count() << " 毫秒]" << std::endl;
        
        // 更新历史
        history = model->MakeHistory(history, round, input, response);
        round++;
        
        std::cout << "----------------------------------------" << std::endl;
    }
    
    return 0;
}