wangjue/llm-ocr
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: Reformat code

Browse Source
This commit is contained in:
wangjue 2026-05-29 16:02:33 +08:00
parent 14bd455bde
commit cf7712830b
1 changed files with 458 additions and 395 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

239
main.cpp
View File

@ -1,36 +1,36 @@
#include "llama.h"
#include "mtmd.h"
#include "mtmd-helper.h"
#include "ggml.h" #include "ggml.h"
#include "llama.h"
#include "mtmd-helper.h"
#include "mtmd.h"
#include <cinttypes>
#include <cmath>
#include <csignal>
#include <cstdio> #include <cstdio>
#include <cstdlib> #include <cstdlib>
#include <cstring> #include <getopt.h>
#include <cmath>
#include <string> #include <string>
#include <vector> #include <sys/stat.h>
#include <thread> #include <thread>
#include <unistd.h> #include <unistd.h>
#include <getopt.h> #include <vector>
#include <signal.h>
#include <inttypes.h>
#include <sys/stat.h>
static volatile bool g_interrupted = false; static volatile bool g_interrupted = false;
static void sigint_handler(int signo) { static void sigint_handler(const int signo) {
(void)signo; (void)signo;
g_interrupted = true; g_interrupted = true;
} }
static void batch_clear(struct llama_batch & batch) { static void batch_clear(llama_batch &batch) { batch.n_tokens = 0; }
batch.n_tokens = 0;
}
static void batch_add(struct llama_batch & batch, llama_token id, llama_pos pos, const std::vector<llama_seq_id> & seq_ids, bool logits) { static void batch_add(llama_batch &batch, const llama_token id,
const llama_pos pos,
const std::vector<llama_seq_id> &seq_ids,
const bool logits) {
batch.token[batch.n_tokens] = id; batch.token[batch.n_tokens] = id;
batch.pos[batch.n_tokens] = pos; batch.pos[batch.n_tokens] = pos;
batch.n_seq_id[batch.n_tokens] = (int32_t)seq_ids.size(); batch.n_seq_id[batch.n_tokens] = static_cast<int32_t>(seq_ids.size());
for (size_t i = 0; i < seq_ids.size(); ++i) { for (size_t i = 0; i < seq_ids.size(); ++i) {
batch.seq_id[batch.n_tokens][i] = seq_ids[i]; batch.seq_id[batch.n_tokens][i] = seq_ids[i];
} }
@ -38,13 +38,16 @@ static void batch_add(struct llama_batch & batch, llama_token id, llama_pos pos,
batch.n_tokens++; batch.n_tokens++;
} }
static std::string token_to_piece(const struct llama_vocab * vocab, llama_token token) { static std::string token_to_piece(const llama_vocab *vocab,
const llama_token token) {
std::string piece; std::string piece;
piece.resize(256); piece.resize(256);
int32_t n = llama_token_to_piece(vocab, token, &piece[0], (int32_t)piece.size(), 0, false); int32_t n = llama_token_to_piece(
vocab, token, &piece[0], static_cast<int32_t>(piece.size()), 0, false);
if (n < 0) { if (n < 0) {
piece.resize(-n); piece.resize(-n);
n = llama_token_to_piece(vocab, token, &piece[0], (int32_t)piece.size(), 0, false); n = llama_token_to_piece(vocab, token, &piece[0],
static_cast<int32_t>(piece.size()), 0, false);
} }
piece.resize(std::max(0, n)); piece.resize(std::max(0, n));
return piece; return piece;
@ -53,21 +56,33 @@ static std::string token_to_piece(const struct llama_vocab * vocab, llama_token
static std::string json_escape(const std::string &s) { static std::string json_escape(const std::string &s) {
std::string out; std::string out;
out.reserve(s.size() + 2); out.reserve(s.size() + 2);
for (char c : s) { for (const char c : s) {
switch (c) { switch (c) {
case '"': out += "\\\""; break; case '"':
case '\\': out += "\\\\"; break; out += "\\\"";
case '\n': out += "\\n"; break; break;
case '\r': out += "\\r"; break; case '\\':
case '\t': out += "\\t"; break; out += "\\\\";
default: out += c; break; break;
case '\n':
out += "\\n";
break;
case '\r':
out += "\\r";
break;
case '\t':
out += "\\t";
break;
default:
out += c;
break;
} }
} }
return out; return out;
} }
static std::vector<unsigned char> read_clipboard() { static std::vector<unsigned char> read_clipboard() {
const char * cmd = nullptr; const char *cmd;
if (getenv("WAYLAND_DISPLAY")) { if (getenv("WAYLAND_DISPLAY")) {
cmd = "wl-paste 2>/dev/null"; cmd = "wl-paste 2>/dev/null";
} else { } else {
@ -75,7 +90,8 @@ static std::vector<unsigned char> read_clipboard() {
} }
std::vector<unsigned char> data; std::vector<unsigned char> data;
FILE *pipe = popen(cmd, "re"); FILE *pipe = popen(cmd, "re");
if (!pipe) return data; if (!pipe)
return data;
unsigned char buf[65536]; unsigned char buf[65536];
size_t n; size_t n;
while ((n = fread(buf, 1, sizeof(buf), pipe)) > 0) { while ((n = fread(buf, 1, sizeof(buf), pipe)) > 0) {
@ -99,29 +115,36 @@ static std::vector<unsigned char> read_clipboard() {
static std::string capture_screenshot() { static std::string capture_screenshot() {
const char *tmpdir = getenv("TMPDIR"); const char *tmpdir = getenv("TMPDIR");
if (!tmpdir) tmpdir = "/tmp"; if (!tmpdir)
tmpdir = "/tmp";
std::string path = std::string(tmpdir) + "/ocr_screenshot.png"; std::string path = std::string(tmpdir) + "/ocr_screenshot.png";
fprintf(stderr, "Select a screen region to OCR...\n"); fprintf(stderr, "Select a screen region to OCR...\n");
auto check = [&]() -> bool { auto check = [&]() -> bool {
struct stat st; struct stat st {};
return stat(path.c_str(), &st) == 0 && st.st_size > 0; return stat(path.c_str(), &st) == 0 && st.st_size > 0;
}; };
// 1. gnome-screenshot (blocks until selection complete) // 1. gnome-screenshot (blocks until selection complete)
unlink(path.c_str()); unlink(path.c_str());
if (system(("gnome-screenshot --area -f " + path + " >/dev/null 2>&1").c_str()) == 0 && check()) if (system(("gnome-screenshot --area -f " + path + " >/dev/null 2>&1")
.c_str()) == 0 &&
check())
return path; return path;
// 2. spectacle (KDE, blocks) // 2. spectacle (KDE, blocks)
unlink(path.c_str()); unlink(path.c_str());
if (system(("spectacle --region -b -o " + path + " >/dev/null 2>&1").c_str()) == 0 && check()) if (system(
("spectacle --region -b -o " + path + " >/dev/null 2>&1").c_str()) ==
0 &&
check())
return path; return path;
// 3. flameshot (writes raw PNG to stdout, blocks) // 3. flameshot (writes raw PNG to stdout, blocks)
unlink(path.c_str()); unlink(path.c_str());
if (system(("flameshot gui -r > " + path + " 2>/dev/null").c_str()) == 0 && check()) if (system(("flameshot gui -r > " + path + " 2>/dev/null").c_str()) == 0 &&
check())
return path; return path;
// 4. maim (lightweight X11, blocks) // 4. maim (lightweight X11, blocks)
@ -131,14 +154,16 @@ static std::string capture_screenshot() {
// 5. slurp + grim (Sway/Wayland, blocks) // 5. slurp + grim (Sway/Wayland, blocks)
unlink(path.c_str()); unlink(path.c_str());
if (system(("slurp | grim -g - " + path + " 2>/dev/null").c_str()) == 0 && check()) if (system(("slurp | grim -g - " + path + " 2>/dev/null").c_str()) == 0 &&
check())
return path; return path;
return ""; return "";
} }
static void print_usage(const char *prog) { static void print_usage(const char *prog) {
fprintf(stderr, fprintf(
stderr,
"Usage: %s -m <model> --mmproj <mmproj> [options] [<image_path>]\n" "Usage: %s -m <model> --mmproj <mmproj> [options] [<image_path>]\n"
"\n" "\n"
"Options:\n" "Options:\n"
@ -161,7 +186,8 @@ static void print_usage(const char * prog) {
"Examples:\n" "Examples:\n"
" %s -m model.gguf --mmproj mmproj.gguf -p \"OCR\" image.png\n" " %s -m model.gguf --mmproj mmproj.gguf -p \"OCR\" image.png\n"
" %s -m model.gguf --mmproj mmproj.gguf # use clipboard image\n" " %s -m model.gguf --mmproj mmproj.gguf # use clipboard image\n"
" %s -m model.gguf --mmproj mmproj.gguf --screenshot # select region\n", " %s -m model.gguf --mmproj mmproj.gguf --screenshot # select "
"region\n",
prog, prog, prog, prog); prog, prog, prog, prog);
} }
@ -172,22 +198,15 @@ int main(int argc, char ** argv) {
std::string prompt = "OCR"; std::string prompt = "OCR";
std::string chat_template; std::string chat_template;
bool screenshot_mode = false; bool screenshot_mode = false;
int n_threads = (int)std::thread::hardware_concurrency(); int n_threads = static_cast<int>(std::thread::hardware_concurrency());
int n_gpu_layers = -1; int n_gpu_layers = -1;
int n_ctx = 8192; int n_ctx = 8192;
float temp = 0.0f; float temp = 0.0f;
uint32_t seed = LLAMA_DEFAULT_SEED; uint32_t seed = LLAMA_DEFAULT_SEED;
bool json_output = false; bool json_output = false;
// pre-process -mm to --mmproj so getopt handles it correctly while (true) {
for (int i = 1; i < argc; i++) { static option long_opts[] = {
if (strcmp(argv[i], "-mm") == 0) {
argv[i] = (char *)"--mmproj";
}
}
while (1) {
static struct option long_opts[] = {
{"temp", required_argument, nullptr, 0}, {"temp", required_argument, nullptr, 0},
{"json", no_argument, nullptr, 1}, {"json", no_argument, nullptr, 1},
{"mmproj", required_argument, nullptr, 2}, {"mmproj", required_argument, nullptr, 2},
@ -197,25 +216,51 @@ int main(int argc, char ** argv) {
{"chat-template", required_argument, nullptr, 4}, {"chat-template", required_argument, nullptr, 4},
{"screenshot", no_argument, nullptr, 5}, {"screenshot", no_argument, nullptr, 5},
{"help", no_argument, nullptr, 'h'}, {"help", no_argument, nullptr, 'h'},
{nullptr, 0, nullptr, 0} {nullptr, 0, nullptr, 0}};
};
int idx = 0; int idx = 0;
int c = getopt_long(argc, argv, "m:p:t:c:s:h", long_opts, &idx); int c = getopt_long(argc, argv, "m:p:t:c:s:h", long_opts, &idx);
if (c == -1) break; if (c == -1)
break;
switch (c) { switch (c) {
case 0: temp = std::stof(optarg); break; case 0:
case 1: json_output = true; break; temp = std::stof(optarg);
case 2: mmproj_path = optarg; break; break;
case 3: n_gpu_layers = std::stoi(optarg); break; case 1:
case 4: chat_template = optarg; break; json_output = true;
case 5: screenshot_mode = true; break; break;
case 'm': model_path = optarg; break; case 2:
case 'p': prompt = optarg; break; mmproj_path = optarg;
case 't': n_threads = std::stoi(optarg); break; break;
case 'c': n_ctx = std::stoi(optarg); break; case 3:
case 's': seed = (uint32_t)std::stoul(optarg); break; n_gpu_layers = std::stoi(optarg);
case 'h': print_usage(argv[0]); return 0; break;
default: print_usage(argv[0]); return 1; case 4:
chat_template = optarg;
break;
case 5:
screenshot_mode = true;
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 = static_cast<uint32_t>(std::stoul(optarg));
break;
case 'h':
print_usage(argv[0]);
return 0;
default:
print_usage(argv[0]);
return 1;
} }
} }
@ -227,7 +272,9 @@ int main(int argc, char ** argv) {
if (screenshot_mode) { if (screenshot_mode) {
image_path = capture_screenshot(); image_path = capture_screenshot();
if (image_path.empty()) { if (image_path.empty()) {
fprintf(stderr, "ERROR: no screenshot tool found. Install gnome-screenshot, flameshot, spectacle, maim, or slurp+grim.\n"); fprintf(stderr,
"ERROR: no screenshot tool found. Install gnome-screenshot, "
"flameshot, spectacle, maim, or slurp+grim.\n");
return 1; return 1;
} }
} else if (optind < argc) { } else if (optind < argc) {
@ -245,12 +292,13 @@ int main(int argc, char ** argv) {
llama_model *model = llama_model_load_from_file(model_path.c_str(), mparams); llama_model *model = llama_model_load_from_file(model_path.c_str(), mparams);
if (!model) { if (!model) {
fprintf(stderr, "ERROR: failed to load model from %s\n", model_path.c_str()); fprintf(stderr, "ERROR: failed to load model from %s\n",
model_path.c_str());
return 1; return 1;
} }
llama_context_params cparams = llama_context_default_params(); llama_context_params cparams = llama_context_default_params();
cparams.n_ctx = (uint32_t)n_ctx; cparams.n_ctx = static_cast<uint32_t>(n_ctx);
cparams.n_batch = 512; cparams.n_batch = 512;
cparams.n_ubatch = 512; cparams.n_ubatch = 512;
@ -263,15 +311,17 @@ int main(int argc, char ** argv) {
llama_set_n_threads(lctx, n_threads, n_threads); llama_set_n_threads(lctx, n_threads, n_threads);
const struct llama_vocab * vocab = llama_model_get_vocab(model); const llama_vocab *vocab = llama_model_get_vocab(model);
mtmd_context_params mtmd_params = mtmd_context_params_default(); mtmd_context_params mtmd_params = mtmd_context_params_default();
mtmd_params.use_gpu = (n_gpu_layers != 0); mtmd_params.use_gpu = (n_gpu_layers != 0);
mtmd_params.n_threads = n_threads; mtmd_params.n_threads = n_threads;
mtmd_context * ctx_vision = mtmd_init_from_file(mmproj_path.c_str(), model, mtmd_params); mtmd_context *ctx_vision =
mtmd_init_from_file(mmproj_path.c_str(), model, mtmd_params);
if (!ctx_vision) { if (!ctx_vision) {
fprintf(stderr, "ERROR: failed to load mmproj from %s\n", mmproj_path.c_str()); fprintf(stderr, "ERROR: failed to load mmproj from %s\n",
mmproj_path.c_str());
llama_free(lctx); llama_free(lctx);
llama_model_free(model); llama_model_free(model);
return 1; return 1;
@ -279,12 +329,14 @@ int main(int argc, char ** argv) {
int64_t t_loaded_us = ggml_time_us(); int64_t t_loaded_us = ggml_time_us();
mtmd_bitmap * bmp = nullptr; mtmd_bitmap *bmp;
if (!image_path.empty()) { if (!image_path.empty()) {
bmp = mtmd_helper_bitmap_init_from_file(ctx_vision, image_path.c_str()); bmp = mtmd_helper_bitmap_init_from_file(ctx_vision, image_path.c_str());
if (!bmp) { if (!bmp) {
fprintf(stderr, "ERROR: failed to load image from %s\n", image_path.c_str()); fprintf(stderr, "ERROR: failed to load image from %s\n",
if (screenshot_mode) unlink(image_path.c_str()); image_path.c_str());
if (screenshot_mode)
unlink(image_path.c_str());
mtmd_free(ctx_vision); mtmd_free(ctx_vision);
llama_free(lctx); llama_free(lctx);
llama_model_free(model); llama_model_free(model);
@ -294,14 +346,16 @@ int main(int argc, char ** argv) {
fprintf(stderr, "Reading image from clipboard...\n"); fprintf(stderr, "Reading image from clipboard...\n");
auto clip = read_clipboard(); auto clip = read_clipboard();
if (clip.empty()) { if (clip.empty()) {
fprintf(stderr, "ERROR: clipboard is empty or no image found.\n" fprintf(stderr,
"ERROR: clipboard is empty or no image found.\n"
"Make sure wl-paste (Wayland) or xclip (X11) is installed.\n"); "Make sure wl-paste (Wayland) or xclip (X11) is installed.\n");
mtmd_free(ctx_vision); mtmd_free(ctx_vision);
llama_free(lctx); llama_free(lctx);
llama_model_free(model); llama_model_free(model);
return 1; return 1;
} }
bmp = mtmd_helper_bitmap_init_from_buf(ctx_vision, clip.data(), clip.size()); bmp =
mtmd_helper_bitmap_init_from_buf(ctx_vision, clip.data(), clip.size());
if (!bmp) { if (!bmp) {
fprintf(stderr, "ERROR: failed to decode clipboard image.\n"); fprintf(stderr, "ERROR: failed to decode clipboard image.\n");
mtmd_free(ctx_vision); mtmd_free(ctx_vision);
@ -330,11 +384,12 @@ int main(int argc, char ** argv) {
// Fallback: use chatml format for models without recognizable template // Fallback: use chatml format for models without recognizable template
if (full_prompt.empty()) { if (full_prompt.empty()) {
full_prompt = "<|im_start|>user\n" + user_content + "<|im_end|>\n<|im_start|>assistant\n"; full_prompt = "<|im_start|>user\n" + user_content +
"<|im_end|>\n<|im_start|>assistant\n";
} }
// Template output already includes role markers and special tokens (e.g. [gMASK]<sop>) // Template output already includes role markers and special tokens (e.g.
// So we don't add BOS separately — add_special=false // [gMASK]<sop>) So we don't add BOS separately — add_special=false
mtmd_input_chunks *chunks = mtmd_input_chunks_init(); mtmd_input_chunks *chunks = mtmd_input_chunks_init();
mtmd_input_text in_text; mtmd_input_text in_text;
@ -349,7 +404,8 @@ int main(int argc, char ** argv) {
fprintf(stderr, "ERROR: mtmd_tokenize failed (%d)\n", res); fprintf(stderr, "ERROR: mtmd_tokenize failed (%d)\n", res);
mtmd_input_chunks_free(chunks); mtmd_input_chunks_free(chunks);
mtmd_bitmap_free(bmp); mtmd_bitmap_free(bmp);
if (screenshot_mode) unlink(image_path.c_str()); if (screenshot_mode)
unlink(image_path.c_str());
mtmd_free(ctx_vision); mtmd_free(ctx_vision);
llama_free(lctx); llama_free(lctx);
llama_model_free(model); llama_model_free(model);
@ -357,12 +413,14 @@ int main(int argc, char ** argv) {
} }
llama_pos n_past = 0; llama_pos n_past = 0;
res = mtmd_helper_eval_chunks(ctx_vision, lctx, chunks, n_past, 0, 512, true, &n_past); res = mtmd_helper_eval_chunks(ctx_vision, lctx, chunks, n_past, 0, 512, true,
&n_past);
if (res != 0) { if (res != 0) {
fprintf(stderr, "ERROR: mtmd_helper_eval_chunks failed (%d)\n", res); fprintf(stderr, "ERROR: mtmd_helper_eval_chunks failed (%d)\n", res);
mtmd_input_chunks_free(chunks); mtmd_input_chunks_free(chunks);
mtmd_bitmap_free(bmp); mtmd_bitmap_free(bmp);
if (screenshot_mode) unlink(image_path.c_str()); if (screenshot_mode)
unlink(image_path.c_str());
mtmd_free(ctx_vision); mtmd_free(ctx_vision);
llama_free(lctx); llama_free(lctx);
llama_model_free(model); llama_model_free(model);
@ -371,15 +429,16 @@ int main(int argc, char ** argv) {
mtmd_input_chunks_free(chunks); mtmd_input_chunks_free(chunks);
mtmd_bitmap_free(bmp); mtmd_bitmap_free(bmp);
if (screenshot_mode) unlink(image_path.c_str()); if (screenshot_mode)
unlink(image_path.c_str());
struct llama_sampler * smpl = nullptr; llama_sampler *smpl;
if (temp <= 0.0f) { if (temp <= 0.0f) {
struct llama_sampler_chain_params sparams = llama_sampler_chain_default_params(); llama_sampler_chain_params sparams = llama_sampler_chain_default_params();
smpl = llama_sampler_chain_init(sparams); smpl = llama_sampler_chain_init(sparams);
llama_sampler_chain_add(smpl, llama_sampler_init_greedy()); llama_sampler_chain_add(smpl, llama_sampler_init_greedy());
} else { } else {
struct llama_sampler_chain_params sparams = llama_sampler_chain_default_params(); llama_sampler_chain_params sparams = llama_sampler_chain_default_params();
smpl = llama_sampler_chain_init(sparams); 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_k(40));
llama_sampler_chain_add(smpl, llama_sampler_init_top_p(0.9f, 1)); llama_sampler_chain_add(smpl, llama_sampler_init_top_p(0.9f, 1));
@ -387,19 +446,21 @@ int main(int argc, char ** argv) {
llama_sampler_chain_add(smpl, llama_sampler_init_dist(seed)); llama_sampler_chain_add(smpl, llama_sampler_init_dist(seed));
} }
struct llama_batch batch = llama_batch_init(1, 0, 1); llama_batch batch = llama_batch_init(1, 0, 1);
int64_t t_infer_start_us = ggml_time_us(); int64_t t_infer_start_us = ggml_time_us();
std::string ocr_text; std::string ocr_text;
int n_generated = 0; int n_generated = 0;
for (int i = 0; i < 2048; i++) { for (int i = 0; i < 2048; i++) {
if (g_interrupted) break; if (g_interrupted)
break;
llama_token token_id = llama_sampler_sample(smpl, lctx, -1); llama_token token_id = llama_sampler_sample(smpl, lctx, -1);
llama_sampler_accept(smpl, token_id); llama_sampler_accept(smpl, token_id);
if (llama_vocab_is_eog(vocab, token_id)) break; if (llama_vocab_is_eog(vocab, token_id))
break;
std::string piece = token_to_piece(vocab, token_id); std::string piece = token_to_piece(vocab, token_id);
ocr_text += piece; ocr_text += piece;
@ -440,11 +501,13 @@ int main(int argc, char ** argv) {
printf(" \"model\": \"%s\",\n", json_escape(model_path).c_str()); printf(" \"model\": \"%s\",\n", json_escape(model_path).c_str());
printf(" \"mmproj\": \"%s\",\n", json_escape(mmproj_path).c_str()); printf(" \"mmproj\": \"%s\",\n", json_escape(mmproj_path).c_str());
printf(" \"prompt\": \"%s\",\n", json_escape(prompt).c_str()); printf(" \"prompt\": \"%s\",\n", json_escape(prompt).c_str());
printf(" \"n_prompt_tokens\": %" PRId64 ",\n", (int64_t)n_past); printf(" \"n_prompt_tokens\": %" PRId64 ",\n",
static_cast<int64_t>(n_past));
printf(" \"n_generated\": %d,\n", n_generated); printf(" \"n_generated\": %d,\n", n_generated);
printf(" \"timings_ms\": {\n"); printf(" \"timings_ms\": {\n");
printf(" \"load\": %" PRId64 ",\n", (t_loaded_us - t_start_us) / 1000); printf(" \"load\": %" PRId64 ",\n", (t_loaded_us - t_start_us) / 1000);
printf(" \"inference\": %" PRId64 "\n", (t_end_us - t_infer_start_us) / 1000); printf(" \"inference\": %" PRId64 "\n",
(t_end_us - t_infer_start_us) / 1000);
printf(" }\n"); printf(" }\n");
printf(" }\n"); printf(" }\n");
printf("}\n"); printf("}\n");