|
| 1 | +// Copyright (c) MLLM Team. |
| 2 | +// Licensed under the MIT License. |
| 3 | + |
| 4 | +#include <fmt/core.h> |
| 5 | + |
| 6 | +#include <mllm/engine/Context.hpp> |
| 7 | +#include <mllm/mllm.hpp> |
| 8 | +#include <mllm/models/ling3/modeling_ling3.hpp> |
| 9 | +#include <mllm/models/ling3/tokenization_ling3.hpp> |
| 10 | +#include <mllm/utils/AnyValue.hpp> |
| 11 | + |
| 12 | +#include <cstdio> |
| 13 | +#include <fstream> |
| 14 | +#include <iostream> |
| 15 | +#include <stdexcept> |
| 16 | +#include <string> |
| 17 | + |
| 18 | +using mllm::Argparse; |
| 19 | + |
| 20 | +namespace { |
| 21 | + |
| 22 | +std::string readPromptFile(const std::string& path) { |
| 23 | + std::ifstream stream(path, std::ios::binary); |
| 24 | + if (!stream) { throw std::invalid_argument("unable to read prompt_file: " + path); } |
| 25 | + std::string text{std::istreambuf_iterator<char>(stream), std::istreambuf_iterator<char>()}; |
| 26 | + while (!text.empty() && (text.back() == '\n' || text.back() == '\r')) { text.pop_back(); } |
| 27 | + if (text.empty()) { throw std::invalid_argument("prompt_file must not be empty: " + path); } |
| 28 | + return text; |
| 29 | +} |
| 30 | + |
| 31 | +} // namespace |
| 32 | + |
| 33 | +MLLM_MAIN({ |
| 34 | + auto engine_args = mllm::engineArgAttach(); |
| 35 | + auto& help = Argparse::add<bool>("-h|--help").help("Show help message"); |
| 36 | + auto& model_path = Argparse::add<std::string>("-m|--model_path").help("MLLM V2 model path").required(true); |
| 37 | + auto& tokenizer_path = Argparse::add<std::string>("-t|--tokenizer_path").help("Official tokenizer.json").required(true); |
| 38 | + auto& config_path = Argparse::add<std::string>("-c|--config_path").help("Ling-3 mobile runtime config").required(true); |
| 39 | + auto& prompt = Argparse::add<std::string>("-p|--prompt").help("Run one prompt non-interactively").required(false); |
| 40 | + auto& prompt_file = Argparse::add<std::string>("--prompt_file").help("Read a prompt from a UTF-8 file").required(false); |
| 41 | + auto& max_new_tokens = |
| 42 | + Argparse::add<int>("-g|--max_new_tokens").help("Maximum generated tokens (default: 8)").required(false); |
| 43 | + auto& min_new_tokens = Argparse::add<int>("--min_new_tokens").help("Suppress EOS until this many tokens").required(false); |
| 44 | + auto& disable_thinking = |
| 45 | + Argparse::add<bool>("--disable_thinking").help("Use the official thinking-off chat template").required(false); |
| 46 | + auto& print_token_ids = Argparse::add<bool>("--print_token_ids").help("Print generated token IDs to stderr").required(false); |
| 47 | + |
| 48 | + for (int index = 1; index < argc; ++index) { |
| 49 | + if (std::string(argv[index]) == "-h" || std::string(argv[index]) == "--help") { |
| 50 | + Argparse::printHelp(); |
| 51 | + return 0; |
| 52 | + } |
| 53 | + } |
| 54 | + Argparse::parse(argc, argv); |
| 55 | + mllm::configEngineWithArgs(engine_args); |
| 56 | + (void)help; |
| 57 | + |
| 58 | + const auto config = mllm::models::ling3::Ling3Config(config_path.get()); |
| 59 | + int generation_limit = max_new_tokens.isSet() ? max_new_tokens.get() : 8; |
| 60 | + int minimum_generation = min_new_tokens.isSet() ? min_new_tokens.get() : 0; |
| 61 | + if (generation_limit <= 0 || generation_limit > config.max_cache_length || minimum_generation < 0 |
| 62 | + || minimum_generation > generation_limit) { |
| 63 | + throw std::invalid_argument("generation lengths must satisfy 0 <= min_new_tokens <= max_new_tokens <= max_cache_length"); |
| 64 | + } |
| 65 | + if (prompt.isSet() && prompt_file.isSet()) { throw std::invalid_argument("prompt and prompt_file are mutually exclusive"); } |
| 66 | + |
| 67 | + std::string configured_prompt; |
| 68 | + if (prompt_file.isSet()) { |
| 69 | + configured_prompt = readPromptFile(prompt_file.get()); |
| 70 | + } else if (prompt.isSet()) { |
| 71 | + configured_prompt = prompt.get(); |
| 72 | + } |
| 73 | + |
| 74 | + const auto parameters = mllm::load(model_path.get(), mllm::ModelFileVersion::kV2); |
| 75 | + mllm::models::ling3::validateLing3ModelConfigMatch(config, parameters); |
| 76 | + auto tokenizer = mllm::models::ling3::Ling3Tokenizer(tokenizer_path.get()); |
| 77 | + auto model = mllm::models::ling3::Ling3ForCausalLM(config); |
| 78 | + model.load(parameters); |
| 79 | + fmt::print("Ling-3.0-tiny: {} layers ({} MLA + {} KDA), CPU threads={}\n", config.num_hidden_layers, |
| 80 | + config.numFullAttentionLayers(), config.numKDALayers(), mllm::Context::instance().getCpuOpThreads()); |
| 81 | + |
| 82 | + int exit_code = 0; |
| 83 | + while (true) { |
| 84 | + std::string prompt_text = configured_prompt; |
| 85 | + if (!prompt.isSet() && !prompt_file.isSet()) { |
| 86 | + fmt::print("Prompt text (or 'exit/quit'): "); |
| 87 | + if (!std::getline(std::cin, prompt_text) || prompt_text == "exit" || prompt_text == "quit") { break; } |
| 88 | + } |
| 89 | + if (prompt_text.empty()) { |
| 90 | + if (prompt.isSet() || prompt_file.isSet()) { throw std::invalid_argument("prompt must not be empty"); } |
| 91 | + continue; |
| 92 | + } |
| 93 | + try { |
| 94 | + auto input = tokenizer.convertMessage({.prompt = prompt_text, |
| 95 | + .system_prompt = "", |
| 96 | + .enable_thinking = !(disable_thinking.isSet() && disable_thinking.get())}); |
| 97 | + const int prompt_tokens = input.at("sequence").shape()[1]; |
| 98 | + if (prompt_tokens + generation_limit - 1 > config.max_cache_length) { |
| 99 | + throw std::invalid_argument("prompt plus generation exceeds max_cache_length"); |
| 100 | + } |
| 101 | + model.resetState(); |
| 102 | + fmt::print("LING3_RUN_START prompt_tokens={} max_new_tokens={} min_new_tokens={}\nResponse: ", prompt_tokens, |
| 103 | + generation_limit, minimum_generation); |
| 104 | + int generated_tokens = 0; |
| 105 | + for (const auto& step : model.chat(input, {{"max_length", mllm::AnyValue(generation_limit)}, |
| 106 | + {"min_new_tokens", mllm::AnyValue(minimum_generation)}, |
| 107 | + {"do_sample", mllm::AnyValue(false)}})) { |
| 108 | + if (print_token_ids.isSet() && print_token_ids.get()) { fmt::print(stderr, "LING3_TOKEN_ID:{}\n", step.cur_token_id); } |
| 109 | + fmt::print("{}", tokenizer.detokenizeBytes(step.cur_token_id)); |
| 110 | + std::fflush(stdout); |
| 111 | + ++generated_tokens; |
| 112 | + } |
| 113 | + fmt::print("\nLING3_RUN_OK prompt_tokens={} generated_tokens={}\n", prompt_tokens, generated_tokens); |
| 114 | + } catch (const std::exception& error) { |
| 115 | + fmt::print(stderr, "LING3_RUN_ERROR:{}\n", error.what()); |
| 116 | + exit_code = 1; |
| 117 | + } |
| 118 | + if (prompt.isSet() || prompt_file.isSet()) { break; } |
| 119 | + } |
| 120 | + |
| 121 | + model.perfSummary(); |
| 122 | + mllm::memoryReport(); |
| 123 | + return exit_code; |
| 124 | +}) |
0 commit comments