Skip to content

Commit 9eeceec

Browse files
asg017asg017
committed
progress
1 parent ae178b8 commit 9eeceec

File tree

5 files changed

+322
-35
lines changed

5 files changed

+322
-35
lines changed

‎TODO‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,3 +4,10 @@
44
./o/llama.cpp/embedr/embedr --version
55
./o/llama.cpp/embedr/embedr
66
```
7+
8+
9+
```
10+
EMBEDR_MODEL_PATH=$PWD/models/mxbai-embed-xsmall-v1-f16.gguf ./o/llama.cpp/embedr/embedr backfill tmp.smol.db nyt_headlines headline
11+
12+
./o/llama.cpp/embedr/embedr backfill tmp.smol.db nyt_headlines headline
13+
```

‎embedr.mk‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
dist/mxbai-embed-xsmall-v1-f16.embedr: ./o/llama.cpp/embedr/embedr
2+
cp $< $@
3+
4+
dist/snowflake-arctic-embed-m-v1.5-f16.embedr: ./o/llama.cpp/embedr/embedr embedr.mk
5+
cp $< $@
6+
echo "-m\nmodels/snowflake-arctic-embed-m-v1.5-f16.gguf\n..." > .args
7+
./o/llamafile/zipalign -j0 \
8+
$@ \
9+
models/snowflake-arctic-embed-m-v1.5-f16.gguf \
10+
.args
11+
rm .args
12+
13+
14+
all: dist/mxbai-embed-xsmall-v1-f16.embedr dist/snowflake-arctic-embed-m-v1.5-f16.embedr

‎llama.cpp/embedr/embedr.c‎

Lines changed: 278 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -7,49 +7,293 @@
77
#include "llama.cpp/embedr/sqlite-lembed.h"
88
#include "llama.cpp/embedr/sqlite-csv.h"
99
#include "llama.cpp/embedr/shell.h"
10-
#include "string.h"
11-
int main(int argc, char ** argv) {
12-
int rc;
13-
sqlite3* db;
14-
sqlite3_stmt* stmt;
15-
rc = sqlite3_auto_extension((void (*)())sqlite3_vec_init);
16-
rc = sqlite3_auto_extension((void (*)())sqlite3_lembed_init);
17-
rc = sqlite3_auto_extension((void (*)())sqlite3_csv_init);
18-
FLAGS_READY = 1;
10+
#include <string.h>
1911

20-
printf(
21-
"llamafile-embed %s, SQLite %s, sqlite-vec=%s, sqlite-lembed=%s\n",
22-
LLAMAFILE_VERSION_STRING,
23-
sqlite3_version,
24-
SQLITE_VEC_VERSION,
25-
SQLITE_LEMBED_VERSION
26-
);
12+
#include <stdlib.h>
13+
#include <assert.h>
14+
#include <time.h>
15+
#include <cosmo.h>
2716

28-
if(argc > 1 && (strcmp(argv[1], "sh") == 0)) {
29-
return mn(argc-1, argv+1);
30-
}
3117

32-
rc = sqlite3_open(":memory:", &db);
33-
if(rc != SQLITE_OK) {
34-
printf("x\n");
35-
return 1;
36-
}
18+
int64_t time_ms(void) {
19+
struct timespec ts;
20+
clock_gettime(CLOCK_MONOTONIC, &ts);
21+
return (int64_t)ts.tv_sec*1000 + (int64_t)ts.tv_nsec/1000000;
22+
}
23+
24+
char * EMBEDR_MODEL = NULL;
3725

38-
rc = sqlite3_prepare_v2(db, "select sqlite_version(), vec_version(), lembed_version()", -1, &stmt, NULL);
39-
if(rc != SQLITE_OK) {
40-
printf("a %s\n", sqlite3_errmsg(db));
41-
return 1;
26+
int embedr_sqlite3_init(sqlite3 * db) {
27+
int rc;
28+
29+
rc = sqlite3_vec_init(db, NULL, NULL); assert(rc == SQLITE_OK);
30+
rc = sqlite3_lembed_init(db, NULL, NULL); assert(rc == SQLITE_OK);
31+
rc = sqlite3_csv_init(db, NULL, NULL); assert(rc == SQLITE_OK);
32+
33+
if(!EMBEDR_MODEL) {
34+
return SQLITE_OK;
35+
}
36+
sqlite3_stmt * stmt;
37+
rc = sqlite3_prepare_v2(db, "insert into temp.lembed_models(model) values (?)", -1, &stmt, NULL);
38+
if(rc != SQLITE_OK) {
39+
assert(rc == SQLITE_OK);
40+
return rc;
41+
}
42+
sqlite3_bind_text(stmt, 1, EMBEDR_MODEL, -1, SQLITE_STATIC);
43+
sqlite3_step(stmt);
44+
rc = sqlite3_finalize(stmt);
45+
assert(rc == SQLITE_OK);
46+
47+
return rc;
48+
}
49+
50+
int register_model(sqlite3 * db, char * modelPath) {
51+
int rc;
52+
sqlite3_stmt * stmt;
53+
rc = sqlite3_prepare_v2(db, "insert into temp.lembed_models(model) values (?) ", -1, &stmt, NULL);
54+
if(rc != SQLITE_OK) {
55+
return rc;
56+
}
57+
sqlite3_bind_text(stmt, 1, modelPath, strlen(modelPath), SQLITE_STATIC);
58+
sqlite3_step(stmt);
59+
rc = sqlite3_finalize(stmt);
60+
if(rc != SQLITE_DONE) {
61+
return rc;
62+
}
63+
return SQLITE_OK;
64+
}
65+
66+
int table_exists(sqlite3 * db, const char * table) {
67+
int rc;
68+
sqlite3_stmt * stmt;
69+
rc = sqlite3_prepare_v2(db, "select ? in (select name from pragma_table_list)", -1, &stmt, NULL);
70+
assert(rc == SQLITE_OK);
71+
sqlite3_bind_text(stmt, 1, table, strlen(table), SQLITE_STATIC);
72+
rc = sqlite3_step(stmt);
73+
assert(rc == SQLITE_ROW);
74+
int result = sqlite3_column_int(stmt, 0);
75+
sqlite3_finalize(stmt);
76+
return result;
77+
}
78+
79+
#define BAR_WIDTH 20
80+
void print_progress_bar(long long nEmbed, long long nTotal, long long elapsed_ms) {
81+
float progress = (float)nEmbed / nTotal;
82+
int bar_fill = (int)(progress * BAR_WIDTH);
83+
84+
long long remaining = nTotal - nEmbed;
85+
float rate = (float)nEmbed / (elapsed_ms / 1000.0);
86+
long long remaining_time = (rate > 0) ? remaining / rate : 0;
87+
88+
printf("\r%3d%%|", (int)(progress * 100));
89+
for (int i = 0; i < BAR_WIDTH; i++) {
90+
if (i < bar_fill)
91+
printf("█");
92+
else
93+
printf(" ");
4294
}
95+
printf("| %lld/%lld [%02lld:%02lld<%02lld:%02lld, %.0f/s]",
96+
nEmbed, nTotal,
97+
elapsed_ms / 1000 / 60, elapsed_ms / 1000 % 60,
98+
remaining_time / 60, remaining_time % 60,
99+
rate);
100+
101+
fflush(stdout);
102+
}
103+
104+
int cmd_backfill(char * dbPath, char * table, char * column) {
105+
int rc;
106+
sqlite3* db;
107+
rc = sqlite3_open(dbPath, &db);
108+
if(rc != SQLITE_OK) {
109+
fprintf(stderr, "could not open database");
110+
return rc;
111+
}
112+
113+
rc = embedr_sqlite3_init(db);
114+
assert(rc == SQLITE_OK);
115+
116+
rc = sqlite3_exec(db, "BEGIN;", NULL, NULL, NULL);
117+
assert(rc == SQLITE_OK);
118+
119+
const char *tableEmbeddings = sqlite3_mprintf("%s_embeddings", table);
120+
assert(tableEmbeddings);
121+
122+
if(!(table_exists(db, tableEmbeddings))) {
123+
const char * zSql = sqlite3_mprintf(
124+
"CREATE TABLE \"%w\"(rowid INTEGER PRIMARY KEY, embedding BLOB);"
125+
"INSERT INTO \"%w\"(rowid) SELECT rowid FROM \"%w\";"
126+
"CREATE INDEX \"idx_%w\" ON \"%w\"(embedding) WHERE embedding IS NULL;",
127+
tableEmbeddings,
128+
tableEmbeddings,
129+
table,
130+
tableEmbeddings,
131+
tableEmbeddings
132+
);
133+
rc = sqlite3_exec(db, zSql, NULL, NULL, NULL);
134+
sqlite3_free((void *) zSql);
135+
assert(rc == SQLITE_OK);
136+
}
137+
138+
139+
int64_t nTotal;
140+
{
141+
sqlite3_stmt * stmt;
142+
const char * zSql = sqlite3_mprintf("SELECT count(*) FROM \"%w\" WHERE embedding IS NULL", tableEmbeddings);
143+
assert(zSql);
144+
rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, NULL);
145+
assert(rc == SQLITE_OK);
43146
rc = sqlite3_step(stmt);
44-
if(rc != SQLITE_ROW) {
45-
printf("b\n");
46-
sqlite3_finalize(stmt);
47-
return 1;
147+
assert(rc == SQLITE_ROW);
148+
nTotal = sqlite3_column_int64(stmt, 0);
149+
sqlite3_finalize(stmt);
150+
}
151+
152+
int64_t nRemaining = nTotal;
153+
154+
155+
sqlite3_stmt * stmt;
156+
const char * zSql = sqlite3_mprintf(
157+
" \
158+
WITH chunk AS ( \
159+
SELECT \
160+
e.rowid, \
161+
lembed(\"%w\") AS embedding \
162+
FROM \"%w\" AS e \
163+
LEFT JOIN \"%w\" AS src ON src.rowid = e.rowid \
164+
WHERE e.embedding IS NULL \
165+
LIMIT ? \
166+
) \
167+
UPDATE \"%w\" AS e \
168+
SET embedding = chunk.embedding \
169+
FROM chunk \
170+
WHERE e.rowid = chunk.rowid \
171+
RETURNING rowid \
172+
",
173+
column,
174+
tableEmbeddings,
175+
table,
176+
tableEmbeddings
177+
);
178+
assert(zSql);
179+
180+
rc = sqlite3_prepare_v2(db, zSql, -1, &stmt, NULL);
181+
sqlite3_free((void *) zSql);
182+
assert(rc == SQLITE_OK);
183+
184+
sqlite3_bind_int(stmt, 1, 16);
185+
186+
int64_t nEmbed = 0;
187+
int64_t t0 = time_ms();
188+
189+
while(1){
190+
sqlite3_reset(stmt);
191+
192+
int nChunkEmbed = 0;
193+
while(1) {
194+
rc = sqlite3_step(stmt);
195+
if(rc == SQLITE_DONE) {
196+
break;
197+
}
198+
assert(rc == SQLITE_ROW);
199+
nChunkEmbed++;
48200
}
49-
printf("%s %s %s\n", sqlite3_column_text(stmt, 0), sqlite3_column_text(stmt, 1), sqlite3_column_text(stmt, 2));
201+
if(nChunkEmbed == 0) {
202+
break;
203+
}
204+
nEmbed += nChunkEmbed;
205+
nRemaining -= nChunkEmbed;
206+
print_progress_bar(nEmbed, nTotal, time_ms() - t0);
207+
}
50208

51-
sqlite3_finalize(stmt);
52209

53210

211+
212+
rc = sqlite3_exec(db, "COMMIT;", NULL, NULL, NULL);
213+
assert(rc == SQLITE_OK);
214+
215+
sqlite3_free((void *) tableEmbeddings);
216+
sqlite3_close(db);
217+
return 0;
218+
}
219+
220+
int cmd_embed(char * source) {
221+
int rc;
222+
sqlite3* db;
223+
sqlite3_stmt * stmt;
224+
225+
rc = sqlite3_open(":memory:", &db);
226+
assert(rc == SQLITE_OK);
227+
228+
rc = embedr_sqlite3_init(db);
229+
assert(rc == SQLITE_OK);
230+
231+
rc = sqlite3_prepare_v2(db, "select vec_to_json(lembed(?))", -1, &stmt, NULL);
232+
assert(rc == SQLITE_OK);
233+
234+
sqlite3_bind_text(stmt, 1, source, strlen(source), SQLITE_STATIC);
235+
236+
rc = sqlite3_step(stmt);
237+
assert(rc == SQLITE_ROW);
238+
239+
printf("%.*s", sqlite3_column_bytes(stmt, 0), sqlite3_column_text(stmt, 0));
240+
241+
sqlite3_finalize(stmt);
242+
sqlite3_close(db);
243+
return 0;
244+
}
245+
246+
247+
int cmd_sh(int argc, char * argv[]) {
248+
return mn(argc, argv);
249+
}
250+
251+
int main(int argc, char ** argv) {
252+
int rc;
253+
sqlite3* db;
254+
sqlite3_stmt* stmt;
255+
256+
FLAG_log_disable = 1;
257+
argc = cosmo_args("/zip/.args", &argv);
258+
FLAGS_READY = 1;
259+
260+
char * modelPath = NULL;
261+
for(int i = 1; i < argc; i++) {
262+
char * arg = argv[i];
263+
if(sqlite3_stricmp(arg, "--model") == 0 || sqlite3_stricmp(arg, "-m") == 0) {
264+
assert(++i <= argc);
265+
EMBEDR_MODEL = argv[i];
266+
}
267+
else if(sqlite3_stricmp(arg, "--version") == 0 || sqlite3_stricmp(arg, "-v") == 0) {
268+
fprintf(stderr,
269+
"llamafile-embed %s, SQLite %s, sqlite-vec=%s, sqlite-lembed=%s\n",
270+
LLAMAFILE_VERSION_STRING,
271+
sqlite3_version,
272+
SQLITE_VEC_VERSION,
273+
SQLITE_LEMBED_VERSION
274+
);
275+
return 0;
276+
}
277+
else if(sqlite3_stricmp(arg, "sh") == 0) {
278+
return cmd_sh(argc-i, argv+i);
279+
}
280+
else if(sqlite3_stricmp(arg, "embed") == 0) {
281+
assert(i + 2 == argc);
282+
return cmd_embed(argv[i+1]);
283+
}
284+
else if(sqlite3_stricmp(arg, "backfill") == 0) {
285+
assert(i + 5 == argc);
286+
char * dbpath = argv[i+1];
287+
char * table = argv[i+2];
288+
char * column = argv[i+3];
289+
return cmd_backfill(dbpath, table, column);
290+
}
291+
else {
292+
printf("Unknown arg %s\n", arg);
293+
return 1;
294+
}
295+
}
296+
54297
return 0;
55298
}
299+

‎llama.cpp/embedr/shell.c‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25607,6 +25607,8 @@ static void open_db(ShellState *p, int openFlags){
2560725607
sqlite3_create_function(p->db, "edit", 2, SQLITE_UTF8, 0,
2560825608
editFunc, 0, 0);
2560925609
#endif
25610+
extern int embedr_sqlite3_init(sqlite3 *);
25611+
embedr_sqlite3_init(p->db);
2561025612

2561125613
if( p->openMode==SHELL_OPEN_ZIPFILE ){
2561225614
char *zSql = sqlite3_mprintf(

‎test.sql‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,24 @@
1+
.timer on
2+
13
insert into temp.lembed_models(model) values ('./models/mxbai-embed-xsmall-v1-f16.gguf');
4+
select vec_length(lembed('yo'));
5+
6+
select :query;
7+
8+
select
9+
nyt_headlines.headline,
10+
vec_distance_cosine(lembed(:query), embedding) as distance
11+
from nyt_headlines_embeddings
12+
left join nyt_headlines on nyt_headlines.rowid = nyt_headlines_embeddings.rowid
13+
order by distance
14+
limit 10;
15+
16+
17+
.exit
18+
219

20+
create virtual table temp.x using csv(filename="fine_food_reviews_1k.csv", header=yes);
321

4-
select lembed('yo');
22+
select *
23+
from temp.x
24+
limit 1;

0 commit comments

Comments
 (0)