Files
2026-07-13 13:33:03 +08:00

129 lines
4.5 KiB
C++

#include <fstream>
#include <sstream>
#include <algorithm>
#include <cmath>
#define MNN_OPEN_TIME_TRACE
#include <MNN/AutoTime.hpp>
#include <MNN/expr/ExprCreator.hpp>
#include <MNN/expr/Module.hpp>
#include "llm/llm.hpp"
#include "llmconfig.hpp"
using namespace MNN::Express;
// #define PRINT_LOSS
static MNN::Express::VARP _CrossEntropy(std::vector<MNN::Express::VARP> inputs, int ignore_index) {
auto shape = _Shape(inputs[0], true), oneV = _Unsqueeze(_Scalar<int>(1), {0}), classes = _Slice(shape, oneV, oneV);
auto mask = _OneHot(inputs[1], classes, _Scalar<float>(1), _Scalar<float>(0), 1);
mask = mask * _Cast<float>(_Unsqueeze(_NotEqual(inputs[1], _Scalar<int>(ignore_index)), {1}));
auto log_prob = inputs[0];
log_prob = _Log(_Softmax(inputs[0], 1));
auto temp = log_prob;
auto output = _ReduceSum(mask * _Negative(temp), {1}, false);
output = _ReduceMean(output);
return output;
}
int main(int argc, const char* argv[]) {
if (argc < 3) {
MNN_PRINT("Usage: ./ppl_eval model/config.json wiki_output max_length\n");
return 0;
}
auto llmPath = argv[1];
auto textPath = argv[2];
int maxLength = -1;
if (argc >= 4) {
maxLength = std::stoi(argv[3]);
}
FUNC_PRINT_ALL(llmPath, s);
FUNC_PRINT_ALL(textPath, s);
std::shared_ptr<MNN::Transformer::Llm> llm(MNN::Transformer::Llm::createLLM(llmPath));
{
AUTOTIME;
llm->set_config("{\"all_logits\":true, \"use_template\":false}");
auto res = llm->load();
if (!res) {
MNN_ERROR("Load LLM error\n");
return 0;
}
}
std::string promptPath = std::string(textPath) + "/prompt.txt";
std::vector<int> inputIds;
{
AUTOTIME;
std::ifstream is(promptPath.c_str());
if (is.fail()) {
MNN_ERROR("Load prompt error\n");
return 0;
}
std::ostringstream os;
os << is.rdbuf();
inputIds = llm->tokenizer_encode(os.str());
}
int ignore_index = -100;
std::shared_ptr<MNN::Express::Module> cross;
{
auto x = _Input({}, NCHW);
auto y = _Input({}, NCHW, halide_type_of<int>());
x->setName("x");
y->setName("y");
auto z = _CrossEntropy({x, y}, ignore_index);
z->setName("z");
auto buffer = Variable::save({z});
cross.reset(Module::load({"x", "y"}, {"z"}, (uint8_t*)buffer.data(), buffer.size()));
}
size_t stride = 512;
size_t contextLength = stride + stride / 2;
std::shared_ptr<MNN::Transformer::LlmConfig> lmConfig(new MNN::Transformer::LlmConfig(llmPath));
if (lmConfig->config_.contains("chunk_limits")) {
contextLength = lmConfig->config_["chunk_limits"][0].get<int>();
stride = (contextLength / 3) * 2;
} else if (lmConfig->config_.contains("chunk")) {
contextLength = lmConfig->config_["chunk"].get<int>();
stride = (contextLength / 3) * 2;
}
FUNC_PRINT(contextLength);
FUNC_PRINT(stride);
auto seqLen = inputIds.size();
if (maxLength > 0) {
seqLen = maxLength;
}
size_t prevEnd = 0;
float lossSum = 0.0f;
int lossNumber = 0;
for (size_t begin = 0; begin < seqLen; begin += stride) {
auto end = std::min(begin + contextLength, seqLen);
std::vector<int> chunkIds(end-begin);
::memcpy(chunkIds.data(), inputIds.data() + begin, chunkIds.size() * sizeof(int));
llm->reset();
auto logits = llm->forward(chunkIds);
logits = MNN::Express::_Squeeze(logits, {0});
auto trgLen = end - prevEnd;
if (prevEnd != 0) {
trgLen += 1;
}
std::vector<int> starts = {(int)(chunkIds.size() - trgLen), 0};
std::vector<int> size = {(int)trgLen-1, -1};
auto startVar = MNN::Express::_Const(starts.data(), {2}, MNN::Express::NCHW, halide_type_of<int>());
auto sizeVar = MNN::Express::_Const(size.data(), {2}, MNN::Express::NCHW, halide_type_of<int>());
logits = MNN::Express::_Slice(logits, startVar, sizeVar);
auto target = _Const(chunkIds.data() + starts[0] + 1, {(int)trgLen - 1}, NCHW, halide_type_of<int>());
auto loss = cross->onForward({logits, target})[0]->readMap<float>()[0];
lossSum+=loss;
lossNumber++;
prevEnd = end;
#ifdef PRINT_LOSS
MNN_PRINT("Compute: %d/%d, loss=%f\n", begin, seqLen, loss);
#endif
if (end == seqLen) {
break;
}
}
MNN_PRINT("Perplexity: %f\n", expf(lossSum / (float)lossNumber));
return 0;
}