From cdcd1f720230397473d6c31d511a4f13c3dc8d1d Mon Sep 17 00:00:00 2001 From: wangjue Date: Fri, 29 May 2026 14:57:37 +0800 Subject: [PATCH] feat: implement OCR tool using llama.cpp multimodal API - Support GLM-OCR model with -m (model) and --mmproj (mmproj) options - Custom prompt support via -p (default: OCR) - Greedy/temperature sampling with --temp flag - JSON output with --json flag - GPU acceleration via -ngl option - Use correct GLM-OCR chat template format for prompt construction --- .gitignore | 8 ++ CMakeLists.txt | 37 ++++++ main.cpp | 322 +++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 367 insertions(+) create mode 100644 .gitignore create mode 100644 CMakeLists.txt create mode 100644 main.cpp diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..5553f05 --- /dev/null +++ b/.gitignore @@ -0,0 +1,8 @@ +.idea +.vscode + +DESIGN.md + +deps +model +build \ No newline at end of file diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..1809055 --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,37 @@ +cmake_minimum_required(VERSION 3.15) +project(ocr VERSION 1.0 LANGUAGES CXX) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) + +set(DEPS_DIR "${CMAKE_SOURCE_DIR}/deps/llama.cpp") + +find_library(LLAMA_LIB llama PATHS "${DEPS_DIR}/lib" NO_DEFAULT_PATH) +find_library(MTMD_LIB mtmd PATHS "${DEPS_DIR}/lib" NO_DEFAULT_PATH) +find_library(GGML_LIB ggml PATHS "${DEPS_DIR}/lib" NO_DEFAULT_PATH) +find_library(GGML_BASE ggml-base PATHS "${DEPS_DIR}/lib" NO_DEFAULT_PATH) +find_library(GGML_CPU ggml-cpu PATHS "${DEPS_DIR}/lib" NO_DEFAULT_PATH) + +if(NOT LLAMA_LIB OR NOT MTMD_LIB) + message(FATAL_ERROR "libraries not found in ${DEPS_DIR}/lib") +endif() + +add_executable(ocr main.cpp) +target_include_directories(ocr PRIVATE "${DEPS_DIR}/include") + +set(CUDA_DIR "${DEPS_DIR}/usr/local/cuda-12.2") +set_target_properties(ocr PROPERTIES + BUILD_RPATH "${DEPS_DIR}/lib;${CUDA_DIR}/targets/x86_64-linux/lib" + INSTALL_RPATH "${DEPS_DIR}/lib;${CUDA_DIR}/targets/x86_64-linux/lib" +) + +target_link_libraries(ocr PRIVATE + ${LLAMA_LIB} + ${MTMD_LIB} + ${GGML_LIB} + ${GGML_BASE} + ${GGML_CPU} + pthread + dl + stdc++fs +) diff --git a/main.cpp b/main.cpp new file mode 100644 index 0000000..7b51af4 --- /dev/null +++ b/main.cpp @@ -0,0 +1,322 @@ +#include "llama.h" +#include "mtmd.h" +#include "mtmd-helper.h" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static volatile bool g_interrupted = false; + +static void sigint_handler(int signo) { + (void)signo; + g_interrupted = true; +} + +static void batch_clear(struct llama_batch & batch) { + batch.n_tokens = 0; +} + +static void batch_add(struct llama_batch & batch, llama_token id, llama_pos pos, const std::vector & seq_ids, bool logits) { + batch.token [batch.n_tokens] = id; + batch.pos [batch.n_tokens] = pos; + batch.n_seq_id[batch.n_tokens] = (int32_t)seq_ids.size(); + for (size_t i = 0; i < seq_ids.size(); ++i) { + batch.seq_id[batch.n_tokens][i] = seq_ids[i]; + } + batch.logits [batch.n_tokens] = logits ? 1 : 0; + batch.n_tokens++; +} + +static std::string token_to_piece(const struct llama_vocab * vocab, llama_token token) { + std::string piece; + piece.resize(256); + int32_t n = llama_token_to_piece(vocab, token, &piece[0], (int32_t)piece.size(), 0, false); + if (n < 0) { + piece.resize(-n); + n = llama_token_to_piece(vocab, token, &piece[0], (int32_t)piece.size(), 0, false); + } + piece.resize(std::max(0, n)); + return piece; +} + +static std::string json_escape(const std::string & s) { + std::string out; + out.reserve(s.size() + 2); + for (char c : s) { + switch (c) { + case '"': out += "\\\""; break; + case '\\': out += "\\\\"; break; + case '\n': out += "\\n"; break; + case '\r': out += "\\r"; break; + case '\t': out += "\\t"; break; + default: out += c; break; + } + } + return out; +} + +static void print_usage(const char * prog) { + fprintf(stderr, + "Usage: %s -m --mmproj [options] \n" + "\n" + "Options:\n" + " -m, --model model file path (GGUF)\n" + " --mmproj,--mm mmproj file path\n" + " -p OCR prompt (default: OCR)\n" + " -t number of threads (default: cpu cores)\n" + " --ngl GPU layers (-1 = all, default: -1)\n" + " -c context size (default: 8192)\n" + " --temp sampling temperature (0 = greedy, default: 0)\n" + " -s random seed\n" + " --json output in JSON format\n" + " -h show this help\n" + "\n" + "Example:\n" + " %s -m model.gguf --mmproj mmproj.gguf -p \"OCR\" image.png\n", + prog, prog); +} + +int main(int argc, char ** argv) { + std::string model_path; + std::string mmproj_path; + std::string image_path; + std::string prompt = "OCR"; + int n_threads = (int)std::thread::hardware_concurrency(); + int n_gpu_layers = -1; + int n_ctx = 8192; + float temp = 0.0f; + uint32_t seed = LLAMA_DEFAULT_SEED; + bool json_output = false; + + // pre-process -mm to --mmproj so getopt handles it correctly + for (int i = 1; i < argc; i++) { + if (strcmp(argv[i], "-mm") == 0) { + argv[i] = (char *)"--mmproj"; + } + } + + while (1) { + static struct option long_opts[] = { + {"temp", required_argument, nullptr, 0}, + {"json", no_argument, nullptr, 1}, + {"mmproj", required_argument, nullptr, 2}, + {"mm", required_argument, nullptr, 2}, + {"ngl", required_argument, nullptr, 3}, + {"model", required_argument, nullptr, 'm'}, + {"help", no_argument, nullptr, 'h'}, + {nullptr, 0, nullptr, 0} + }; + int idx = 0; + int c = getopt_long(argc, argv, "m:p:t:c:s:h", long_opts, &idx); + if (c == -1) break; + switch (c) { + case 0: temp = std::stof(optarg); break; + case 1: json_output = true; break; + case 2: mmproj_path = optarg; break; + case 3: n_gpu_layers = std::stoi(optarg); break; + case 'm': model_path = optarg; break; + case 'p': prompt = optarg; break; + case 't': n_threads = std::stoi(optarg); break; + case 'c': n_ctx = std::stoi(optarg); break; + case 's': seed = (uint32_t)std::stoul(optarg); break; + case 'h': print_usage(argv[0]); return 0; + default: print_usage(argv[0]); return 1; + } + } + + if ((optind + 1) != argc) { + print_usage(argv[0]); + return 1; + } + image_path = argv[optind]; + + if (model_path.empty() || mmproj_path.empty()) { + fprintf(stderr, "ERROR: -m/--model and --mmproj/--mm are required\n"); + return 1; + } + + signal(SIGINT, sigint_handler); + + int64_t t_start_us = ggml_time_us(); + + llama_backend_init(); + + llama_model_params mparams = llama_model_default_params(); + mparams.n_gpu_layers = n_gpu_layers; + + llama_model * model = llama_model_load_from_file(model_path.c_str(), mparams); + if (!model) { + fprintf(stderr, "ERROR: failed to load model from %s\n", model_path.c_str()); + return 1; + } + + llama_context_params cparams = llama_context_default_params(); + cparams.n_ctx = (uint32_t)n_ctx; + cparams.n_batch = 512; + cparams.n_ubatch = 512; + + llama_context * lctx = llama_init_from_model(model, cparams); + if (!lctx) { + fprintf(stderr, "ERROR: failed to create context\n"); + llama_model_free(model); + return 1; + } + + llama_set_n_threads(lctx, n_threads, n_threads); + + const struct llama_vocab * vocab = llama_model_get_vocab(model); + + mtmd_context_params mtmd_params = mtmd_context_params_default(); + mtmd_params.use_gpu = (n_gpu_layers != 0); + mtmd_params.n_threads = n_threads; + + mtmd_context * ctx_vision = mtmd_init_from_file(mmproj_path.c_str(), model, mtmd_params); + if (!ctx_vision) { + fprintf(stderr, "ERROR: failed to load mmproj from %s\n", mmproj_path.c_str()); + llama_free(lctx); + llama_model_free(model); + return 1; + } + + int64_t t_loaded_us = ggml_time_us(); + + mtmd_bitmap * bmp = mtmd_helper_bitmap_init_from_file(ctx_vision, image_path.c_str()); + if (!bmp) { + fprintf(stderr, "ERROR: failed to load image from %s\n", image_path.c_str()); + mtmd_free(ctx_vision); + llama_free(lctx); + llama_model_free(model); + return 1; + } + + // Construct prompt using GLM-OCR chat template format + // The model expects: [gMASK]<|user|>\n<__media__>PROMPT\n<|assistant|>\n + std::string full_prompt = "[gMASK]<|user|>\n"; + full_prompt += mtmd_default_marker(); + full_prompt += prompt; + full_prompt += "\n<|assistant|>\n"; + + mtmd_input_chunks * chunks = mtmd_input_chunks_init(); + + mtmd_input_text in_text; + in_text.text = full_prompt.c_str(); + in_text.add_special = false; // template already includes [gMASK]/ + in_text.parse_special = true; + + const mtmd_bitmap * bitmaps[1] = { bmp }; + + int32_t res = mtmd_tokenize(ctx_vision, chunks, &in_text, bitmaps, 1); + if (res != 0) { + fprintf(stderr, "ERROR: mtmd_tokenize failed (%d)\n", res); + mtmd_input_chunks_free(chunks); + mtmd_bitmap_free(bmp); + mtmd_free(ctx_vision); + llama_free(lctx); + llama_model_free(model); + return 1; + } + + llama_pos n_past = 0; + res = mtmd_helper_eval_chunks(ctx_vision, lctx, chunks, n_past, 0, 512, true, &n_past); + if (res != 0) { + fprintf(stderr, "ERROR: mtmd_helper_eval_chunks failed (%d)\n", res); + mtmd_input_chunks_free(chunks); + mtmd_bitmap_free(bmp); + mtmd_free(ctx_vision); + llama_free(lctx); + llama_model_free(model); + return 1; + } + + mtmd_input_chunks_free(chunks); + mtmd_bitmap_free(bmp); + + struct llama_sampler * smpl = nullptr; + if (temp <= 0.0f) { + struct llama_sampler_chain_params sparams = llama_sampler_chain_default_params(); + smpl = llama_sampler_chain_init(sparams); + llama_sampler_chain_add(smpl, llama_sampler_init_greedy()); + } else { + struct llama_sampler_chain_params sparams = llama_sampler_chain_default_params(); + smpl = llama_sampler_chain_init(sparams); + llama_sampler_chain_add(smpl, llama_sampler_init_top_k(40)); + llama_sampler_chain_add(smpl, llama_sampler_init_top_p(0.9f, 1)); + llama_sampler_chain_add(smpl, llama_sampler_init_temp(temp)); + llama_sampler_chain_add(smpl, llama_sampler_init_dist(seed)); + } + + struct llama_batch batch = llama_batch_init(1, 0, 1); + + int64_t t_infer_start_us = ggml_time_us(); + std::string ocr_text; + int n_generated = 0; + + for (int i = 0; i < 2048; i++) { + if (g_interrupted) break; + + llama_token token_id = llama_sampler_sample(smpl, lctx, -1); + llama_sampler_accept(smpl, token_id); + + if (llama_vocab_is_eog(vocab, token_id)) break; + + std::string piece = token_to_piece(vocab, token_id); + ocr_text += piece; + + if (!json_output) { + printf("%s", piece.c_str()); + fflush(stdout); + } + + batch_clear(batch); + batch_add(batch, token_id, n_past++, {0}, true); + + if (llama_decode(lctx, batch) != 0) { + fprintf(stderr, "\nERROR: llama_decode failed\n"); + break; + } + + n_generated++; + } + + int64_t t_end_us = ggml_time_us(); + + if (!json_output) { + printf("\n"); + } + + llama_batch_free(batch); + llama_sampler_free(smpl); + mtmd_free(ctx_vision); + llama_free(lctx); + llama_model_free(model); + llama_backend_free(); + + if (json_output) { + printf("{\n"); + printf(" \"text\": \"%s\",\n", json_escape(ocr_text).c_str()); + printf(" \"meta\": {\n"); + printf(" \"model\": \"%s\",\n", json_escape(model_path).c_str()); + printf(" \"mmproj\": \"%s\",\n", json_escape(mmproj_path).c_str()); + printf(" \"prompt\": \"%s\",\n", json_escape(prompt).c_str()); + printf(" \"n_prompt_tokens\": %" PRId64 ",\n", (int64_t)n_past); + printf(" \"n_generated\": %d,\n", n_generated); + printf(" \"timings_ms\": {\n"); + printf(" \"load\": %" PRId64 ",\n", (t_loaded_us - t_start_us) / 1000); + printf(" \"inference\": %" PRId64 "\n", (t_end_us - t_infer_start_us) / 1000); + printf(" }\n"); + printf(" }\n"); + printf("}\n"); + } + + return g_interrupted ? 130 : 0; +}