Method build

src/framework/sampling/decode_modules.cpp:43–58

}

engine::core::TensorValue GreedyDecodeModule::build(engine::core::ModuleBuildContext & ctx, const engine::core::TensorValue & logits) const {
    if (ctx.ggml == nullptr) {
        throw std::runtime_error("ModuleBuildContext.ggml is null");
    }
    engine::core::validate_rank_between(logits, 3, 3, "logits");

    auto flat = engine::core::reshape_tensor(
        ctx,
        engine::core::ensure_backend_addressable_layout(ctx, logits),
        engine::core::TensorShape::from_dims({logits.shape.num_elements() / logits.shape.last_dim(), logits.shape.last_dim()}));
    auto argmax = engine::core::wrap_tensor(
        ggml_argmax(ctx.ggml, flat.tensor),
        engine::core::TensorShape::from_dims({flat.shape.dims[0]}),
        GGML_TYPE_I32);
    return engine::core::reshape_tensor(ctx, argmax, engine::core::TensorShape::from_dims({logits.shape.dims[0], logits.shape.dims[1]}));
}

const engine::core::ModuleSchema & GreedyDecodeModule::static_schema() noexcept {
    return kGreedyDecodeSchema;