1#include "llama.h"
2#include "vectordb.h"
3#include "models.h"
4
5#define NONSTD_IMPLEMENTATION
6#include "nonstd.h"
7
8#include <stdio.h>
9#include <stdlib.h>
10#include <string.h>
11#include <getopt.h>
12
13static void llama_log_callback(enum ggml_log_level level, const char *text, void *user_data) {
14 (void)level;
15 (void)user_data;
16 (void)text;
17}
18
19static void list_available_models() {
20 printf("Model list:\n");
21 ModelConfig model;
22 static_foreach(ModelConfig, model, models) {
23 printf(" - %s [ctx: %d, temp: %f]\n", model.name, model.n_ctx, model.temperature);
24 }
25}
26
27static void show_help(const char *prog) {
28 printf("Usage: %s [OPTIONS]\n", prog);
29 printf("Options:\n");
30 printf(" -m, --model <name> Specify model to use (default: first model)\n");
31 printf(" -i, --in <file> Specify input context file\n");
32 printf(" -o, --out <file> Specify output vector database file\n");
33 printf(" -l, --list Lists all available models\n");
34 printf(" -v, --verbose Enable verbose logging\n");
35 printf(" -h, --help Show this help message\n");
36}
37
38int main(int argc, char **argv) {
39 set_log_level(LOG_DEBUG);
40
41 const char *model_name = NULL;
42 const char *in_file = NULL;
43 const char *out_file = NULL;
44 int list_models = 0;
45 int verbose = 0;
46
47 static struct option long_options[] = {
48 {"model", required_argument, 0, 'm'},
49 {"in", required_argument, 0, 'i'},
50 {"out", required_argument, 0, 'o'},
51 {"list", no_argument, 0, 'l'},
52 {"verbose", no_argument, 0, 'v'},
53 {"help", no_argument, 0, 'h'},
54 {0, 0, 0, 0}
55 };
56
57 int opt;
58 int option_index = 0;
59 while ((opt = getopt_long(argc, argv, "m:i:o:lvh", long_options, &option_index)) != -1) {
60 switch (opt) {
61 case 'm':
62 model_name = optarg;
63 break;
64 case 'i':
65 in_file = optarg;
66 break;
67 case 'o':
68 out_file = optarg;
69 break;
70 case 'l':
71 list_models = 1;
72 break;
73 case 'v':
74 verbose = 1;
75 break;
76 case 'h':
77 show_help(argv[0]);
78 return 0;
79 default:
80 fprintf(stderr, "Usage: %s [-m model] [-i file] [-o file] [-lvh]\n", argv[0]);
81 return 1;
82 }
83 }
84
85 if (verbose == 0) {
86 llama_log_set(llama_log_callback, NULL);
87 }
88
89 if (list_models == 1) {
90 list_available_models();
91 return 0;
92 }
93
94 if (in_file == NULL) {
95 log_message(stderr, LOG_ERROR, "Input context file must be provided. Exiting...");
96 return 1;
97 }
98
99 if (out_file == NULL) {
100 log_message(stderr, LOG_ERROR, "Output vector context file must be provided. Exiting...");
101 return 1;
102 }
103
104 llama_backend_init();
105
106 const ModelConfig *cfg = NULL;
107 if (model_name != NULL) {
108 cfg = get_model_by_name(model_name);
109 if (cfg == NULL) {
110 log_message(stderr, LOG_ERROR, "Unknown model '%s'", model_name);
111 llama_backend_free();
112 return 1;
113 }
114 } else {
115 cfg = &models[0];
116 }
117
118 struct llama_model_params model_params = llama_model_default_params();
119 model_params.n_gpu_layers = cfg->n_gpu_layers;
120 model_params.use_mmap = cfg->use_mmap;
121 struct llama_model *model = llama_model_load_from_file(cfg->filepath, model_params);
122 if (model == NULL) {
123 log_message(stderr, LOG_ERROR, "Unable to load embedding model");
124 llama_backend_free();
125 return 1;
126 }
127
128 struct llama_context_params cparams = llama_context_default_params();
129 cparams.n_ctx = cfg->n_ctx;
130 cparams.n_batch = cfg->n_batch;
131 cparams.embeddings = true;
132
133 struct llama_context *embed_ctx = llama_init_from_model(model, cparams);
134 if (embed_ctx == NULL) {
135 log_message(stderr, LOG_ERROR, "Failed to create embedding context");
136 llama_model_free(model);
137 llama_backend_free();
138 return 1;
139 }
140
141 FILE *context_fp = fopen(in_file, "r");
142 if (context_fp == NULL) {
143 log_message(stderr, LOG_ERROR, "Unable to open context file %s", in_file);
144 return 1;
145 }
146
147 VectorDB db;
148 vdb_init(&db, embed_ctx);
149
150 char line[1024];
151 while (fgets(line, sizeof(line), context_fp) != NULL) {
152 size_t len = strlen(line);
153 while (len > 0 && (line[len - 1] == '\n' || line[len - 1] == '\r')) {
154 line[len - 1] = '\0';
155 len--;
156 }
157 if (len == 0) {
158 continue;
159 }
160 vdb_add_document(&db, line);
161 }
162
163 VectorDBErrorCode vdb_rc = vdb_save(&db, out_file);
164 if (vdb_rc != VDB_SUCCESS) {
165 log_message(stderr, LOG_ERROR, "Something went wrong saving file %s: %s", out_file, vdb_error(vdb_rc));
166 fclose(context_fp);
167 return 1;
168 }
169
170 log_message(stdout, LOG_INFO, "Context vector database file %s successfully written", out_file);
171 fclose(context_fp);
172 return 0;
173}