refactor: Reformat code
This commit is contained in:
wangjue
parent
14bd455bde
commit
cf7712830b
239
main.cpp
239
main.cpp
@ -1,36 +1,36 @@
|
||||
#include "llama.h"
|
||||
#include "mtmd.h"
|
||||
#include "mtmd-helper.h"
|
||||
#include "ggml.h"
|
||||
#include "llama.h"
|
||||
#include "mtmd-helper.h"
|
||||
#include "mtmd.h"
|
||||
|
||||
#include <cinttypes>
|
||||
#include <cmath>
|
||||
#include <csignal>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <cmath>
|
||||
#include <getopt.h>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <sys/stat.h>
|
||||
#include <thread>
|
||||
#include <unistd.h>
|
||||
#include <getopt.h>
|
||||
#include <signal.h>
|
||||
#include <inttypes.h>
|
||||
#include <sys/stat.h>
|
||||
#include <vector>
|
||||
|
||||
static volatile bool g_interrupted = false;
|
||||
|
||||
static void sigint_handler(int signo) {
|
||||
static void sigint_handler(const int signo) {
|
||||
(void)signo;
|
||||
g_interrupted = true;
|
||||
}
|
||||
|
||||
static void batch_clear(struct llama_batch & batch) {
|
||||
batch.n_tokens = 0;
|
||||
}
|
||||
static void batch_clear(llama_batch &batch) { 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.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) {
|
||||
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++;
|
||||
}
|
||||
|
||||
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;
|
||||
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) {
|
||||
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));
|
||||
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) {
|
||||
std::string out;
|
||||
out.reserve(s.size() + 2);
|
||||
for (char c : s) {
|
||||
for (const 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;
|
||||
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 std::vector<unsigned char> read_clipboard() {
|
||||
const char * cmd = nullptr;
|
||||
const char *cmd;
|
||||
if (getenv("WAYLAND_DISPLAY")) {
|
||||
cmd = "wl-paste 2>/dev/null";
|
||||
} else {
|
||||
@ -75,7 +90,8 @@ static std::vector<unsigned char> read_clipboard() {
|
||||
}
|
||||
std::vector<unsigned char> data;
|
||||
FILE *pipe = popen(cmd, "re");
|
||||
if (!pipe) return data;
|
||||
if (!pipe)
|
||||
return data;
|
||||
unsigned char buf[65536];
|
||||
size_t n;
|
||||
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() {
|
||||
const char *tmpdir = getenv("TMPDIR");
|
||||
if (!tmpdir) tmpdir = "/tmp";
|
||||
if (!tmpdir)
|
||||
tmpdir = "/tmp";
|
||||
std::string path = std::string(tmpdir) + "/ocr_screenshot.png";
|
||||
|
||||
fprintf(stderr, "Select a screen region to OCR...\n");
|
||||
|
||||
auto check = [&]() -> bool {
|
||||
struct stat st;
|
||||
struct stat st {};
|
||||
return stat(path.c_str(), &st) == 0 && st.st_size > 0;
|
||||
};
|
||||
|
||||
// 1. gnome-screenshot (blocks until selection complete)
|
||||
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;
|
||||
|
||||
// 2. spectacle (KDE, blocks)
|
||||
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;
|
||||
|
||||
// 3. flameshot (writes raw PNG to stdout, blocks)
|
||||
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;
|
||||
|
||||
// 4. maim (lightweight X11, blocks)
|
||||
@ -131,14 +154,16 @@ static std::string capture_screenshot() {
|
||||
|
||||
// 5. slurp + grim (Sway/Wayland, blocks)
|
||||
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 "";
|
||||
}
|
||||
|
||||
static void print_usage(const char *prog) {
|
||||
fprintf(stderr,
|
||||
fprintf(
|
||||
stderr,
|
||||
"Usage: %s -m <model> --mmproj <mmproj> [options] [<image_path>]\n"
|
||||
"\n"
|
||||
"Options:\n"
|
||||
@ -161,7 +186,8 @@ static void print_usage(const char * prog) {
|
||||
"Examples:\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 --screenshot # select region\n",
|
||||
" %s -m model.gguf --mmproj mmproj.gguf --screenshot # select "
|
||||
"region\n",
|
||||
prog, prog, prog, prog);
|
||||
}
|
||||
|
||||
@ -172,22 +198,15 @@ int main(int argc, char ** argv) {
|
||||
std::string prompt = "OCR";
|
||||
std::string chat_template;
|
||||
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_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[] = {
|
||||
while (true) {
|
||||
static option long_opts[] = {
|
||||
{"temp", required_argument, nullptr, 0},
|
||||
{"json", no_argument, nullptr, 1},
|
||||
{"mmproj", required_argument, nullptr, 2},
|
||||
@ -197,25 +216,51 @@ int main(int argc, char ** argv) {
|
||||
{"chat-template", required_argument, nullptr, 4},
|
||||
{"screenshot", no_argument, nullptr, 5},
|
||||
{"help", no_argument, nullptr, 'h'},
|
||||
{nullptr, 0, nullptr, 0}
|
||||
};
|
||||
{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;
|
||||
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 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 = (uint32_t)std::stoul(optarg); break;
|
||||
case 'h': print_usage(argv[0]); return 0;
|
||||
default: print_usage(argv[0]); return 1;
|
||||
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 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) {
|
||||
image_path = capture_screenshot();
|
||||
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;
|
||||
}
|
||||
} 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);
|
||||
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;
|
||||
}
|
||||
|
||||
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_ubatch = 512;
|
||||
|
||||
@ -263,15 +311,17 @@ int main(int argc, char ** argv) {
|
||||
|
||||
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_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);
|
||||
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());
|
||||
fprintf(stderr, "ERROR: failed to load mmproj from %s\n",
|
||||
mmproj_path.c_str());
|
||||
llama_free(lctx);
|
||||
llama_model_free(model);
|
||||
return 1;
|
||||
@ -279,12 +329,14 @@ int main(int argc, char ** argv) {
|
||||
|
||||
int64_t t_loaded_us = ggml_time_us();
|
||||
|
||||
mtmd_bitmap * bmp = nullptr;
|
||||
mtmd_bitmap *bmp;
|
||||
if (!image_path.empty()) {
|
||||
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());
|
||||
if (screenshot_mode) unlink(image_path.c_str());
|
||||
fprintf(stderr, "ERROR: failed to load image from %s\n",
|
||||
image_path.c_str());
|
||||
if (screenshot_mode)
|
||||
unlink(image_path.c_str());
|
||||
mtmd_free(ctx_vision);
|
||||
llama_free(lctx);
|
||||
llama_model_free(model);
|
||||
@ -294,14 +346,16 @@ int main(int argc, char ** argv) {
|
||||
fprintf(stderr, "Reading image from clipboard...\n");
|
||||
auto clip = read_clipboard();
|
||||
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");
|
||||
mtmd_free(ctx_vision);
|
||||
llama_free(lctx);
|
||||
llama_model_free(model);
|
||||
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) {
|
||||
fprintf(stderr, "ERROR: failed to decode clipboard image.\n");
|
||||
mtmd_free(ctx_vision);
|
||||
@ -330,11 +384,12 @@ int main(int argc, char ** argv) {
|
||||
|
||||
// Fallback: use chatml format for models without recognizable template
|
||||
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>)
|
||||
// So we don't add BOS separately — add_special=false
|
||||
// Template output already includes role markers and special tokens (e.g.
|
||||
// [gMASK]<sop>) So we don't add BOS separately — add_special=false
|
||||
mtmd_input_chunks *chunks = mtmd_input_chunks_init();
|
||||
|
||||
mtmd_input_text in_text;
|
||||
@ -349,7 +404,8 @@ int main(int argc, char ** argv) {
|
||||
fprintf(stderr, "ERROR: mtmd_tokenize failed (%d)\n", res);
|
||||
mtmd_input_chunks_free(chunks);
|
||||
mtmd_bitmap_free(bmp);
|
||||
if (screenshot_mode) unlink(image_path.c_str());
|
||||
if (screenshot_mode)
|
||||
unlink(image_path.c_str());
|
||||
mtmd_free(ctx_vision);
|
||||
llama_free(lctx);
|
||||
llama_model_free(model);
|
||||
@ -357,12 +413,14 @@ int main(int argc, char ** argv) {
|
||||
}
|
||||
|
||||
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) {
|
||||
fprintf(stderr, "ERROR: mtmd_helper_eval_chunks failed (%d)\n", res);
|
||||
mtmd_input_chunks_free(chunks);
|
||||
mtmd_bitmap_free(bmp);
|
||||
if (screenshot_mode) unlink(image_path.c_str());
|
||||
if (screenshot_mode)
|
||||
unlink(image_path.c_str());
|
||||
mtmd_free(ctx_vision);
|
||||
llama_free(lctx);
|
||||
llama_model_free(model);
|
||||
@ -371,15 +429,16 @@ int main(int argc, char ** argv) {
|
||||
|
||||
mtmd_input_chunks_free(chunks);
|
||||
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) {
|
||||
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);
|
||||
llama_sampler_chain_add(smpl, llama_sampler_init_greedy());
|
||||
} 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);
|
||||
llama_sampler_chain_add(smpl, llama_sampler_init_top_k(40));
|
||||
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));
|
||||
}
|
||||
|
||||
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();
|
||||
std::string ocr_text;
|
||||
int n_generated = 0;
|
||||
|
||||
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_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);
|
||||
ocr_text += piece;
|
||||
@ -440,11 +501,13 @@ int main(int argc, char ** argv) {
|
||||
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_prompt_tokens\": %" PRId64 ",\n",
|
||||
static_cast<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(" \"inference\": %" PRId64 "\n",
|
||||
(t_end_us - t_infer_start_us) / 1000);
|
||||
printf(" }\n");
|
||||
printf(" }\n");
|
||||
printf("}\n");
|
||||
|
||||
Loading…
Reference in New Issue
Block a user