diff --git a/generate_embeddings.lua b/generate_embeddings.lua index d242090..1dcdd03 100755 --- a/generate_embeddings.lua +++ b/generate_embeddings.lua @@ -2,83 +2,21 @@ package.path = (arg[0]:match("@?(.*/)") or arg[0]:match("@?(.*\\)")) .. "lib" .. package.config:sub(1, 1) .. "?.lua;" .. package.path local utility = require "utility" -local json = utility.require("dkjson") -local PATH = "PRIVATE_DATA/notebook" -local whitelist = { md = true, } +local text_processing = utility.require("text_processing") +local timing = utility.require("timing") +local json = utility.require("dkjson") -- local embedding_model = "nomic-embed-text" -- local maximum_file_size = 2048 local embedding_model = "qwen3-embedding:0.6b" local maximum_file_size = 32768 +local NOTEBOOK_PATH = "PRIVATE_DATA/notebook" +local blacklist = utility.enumerate{ ".git", } -- accelerate processing by ignoring .git directory +local extension_whitelist = utility.enumerate{ "md", } -local timing = {} -local function display_timing(n) - local function _display(n) - local time = timing[n] - local previous = timing[n - 1] - - local delta = time.time - previous.time - if delta >= 2*60*60 then -- 2 hours - delta = tostring(math.floor(delta/(60*60/10))/10) .. " hours" - elseif delta >= 120 then -- 2 minutes - delta = tostring(math.floor(delta/(60/10))/10) .. " minutes" - else - delta = tostring(delta) .. " seconds" - end - - print(delta, previous.label) - end - - if n then - _display(n) - else - print("All measured timings:") - for i = 2, #timing do - _display(i) - end - end -end - -local function mark_timing(label) - timing[#timing + 1] = { label = label, time = os.time(), } - if #timing > 1 then - display_timing(#timing) - end -end - --- strip YAML frontmatter (if present) --- can error, will return nil & error message -local function strip_frontmatter(text) - local tab = text:split("\n") - if tab[1] == "---" then - table.remove(tab, 1) - while true do - local done = tab[1] == "---" - table.remove(tab, 1) - if done then - return table.concat(tab, "\n") - elseif #tab < 1 then - return nil, "Invalid YAML frontmatter." - end - end - end - return text -end - -local tree -tree = function(path, fn) - utility.list(path or ".", function(path_name) - if path_name == ".git" then return end - if utility.is_file(path_name) then - fn(path_name) - else - tree(path .. utility.path_separator .. path_name, fn) - end - end) -end local function leftpad(text, length) return string.rep("0", length - #(tostring(text))) .. text @@ -89,36 +27,22 @@ end local file_list = {} local embeddings = {} -mark_timing("Assembling file list.") -tree(PATH, function(file_name) - local path, name, extension = utility.split_path_components(file_name) - if not extension then - return - end - - if not whitelist[extension] then - return - end - +timing.mark("Assembling file list.") +utility.tree(NOTEBOOK_PATH, { + blacklist = blacklist, + extension_whitelist = extension_whitelist, +}, function(file_name) file_list[#file_list + 1] = file_name end) -mark_timing("Generating embeddings.") +timing.mark("Generating embeddings.") for i = 1, #file_list do local function _run() local file_name = file_list[i] -- stripping YAML frontmatter before generating embeddings - local file_contents = utility.open(file_name, "r", function(file) - return file:read("*all") - end) - - local tmp_file_contents, error_message = strip_frontmatter(file_contents) - if tmp_file_contents == nil then - print("ERROR: " .. file_name .. " " .. error_message) - tmp_file_contents = file_contents - end - + local file_contents = utility.read_file(file_name) + local tmp_file_contents = text_processing.strip_frontmatter(file_contents) if #tmp_file_contents == 0 then print(file_name .. "\n is empty and will be skipped.") return @@ -129,14 +53,7 @@ for i = 1, #file_list do tmp_file_contents = tmp_file_contents:sub(1, maximum_file_size) end - local tmp_file_name = utility.tmp_file_name() - utility.open(tmp_file_name, "w", function(file) - file:write(tmp_file_contents) - end) - - local output = utility.capture_safe("cat " .. tmp_file_name:enquote() .. " | ollama run " .. embedding_model) - os.execute("rm " .. tmp_file_name) - output = output:sub(1, -2) -- strip extra newline from utility.capture_safe + local output = utility.llm_prompt(tmp_file_contents, embedding_model) if #output == 0 then print("Warning: " .. file_name .. " did not have embeddings generated.") @@ -157,13 +74,9 @@ for i = 1, #file_list do _run() end -mark_timing("Outputting embeddings in JSON.") -utility.open("PRIVATE_DATA/embeddings.json", "w", function(file) - local output = json.encode(embeddings, { indent = true }) - file:write(output) - file:write("\n") -end) +timing.mark("Outputting embeddings in JSON.") +utility.save_data(embeddings, "PRIVATE_DATA/embeddings.json") -mark_timing("Finished.") +timing.mark("Finished.") print("") -display_timing() +timing.display() diff --git a/lib/timing.lua b/lib/timing.lua new file mode 100644 index 0000000..d272c1f --- /dev/null +++ b/lib/timing.lua @@ -0,0 +1,37 @@ +local timing = {} + +timing.display = function(n) + local function _display(n) + local time = timing[n] + local previous = timing[n - 1] + + local delta = time.time - previous.time + if delta >= 2*60*60 then -- if more than 2 hours + delta = tostring(math.floor(delta/(60*60/10))/10) .. " hours" + elseif delta >= 120 then -- if more then 2 minutes + delta = tostring(math.floor(delta/(60/10))/10) .. " minutes" + else + delta = tostring(delta) .. " seconds" + end + + print(delta, previous.label) + end + + if n then + _display(n) + else + print("All measured timings:") + for i = 2, #timing do + _display(i) + end + end +end + +timing.mark = function(label) + timing[#timing + 1] = { label = label, time = os.time(), } + if #timing > 1 then + timing.display(#timing) + end +end + +return timing