diff --git a/generate-embeddings.lua b/generate-embeddings.lua index baf7397..b759000 100755 --- a/generate-embeddings.lua +++ b/generate-embeddings.lua @@ -7,6 +7,25 @@ local json = utility.require("dkjson") local PATH = "notebook" local whitelist = { md = true, } +-- 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) @@ -18,6 +37,8 @@ tree = function(path, fn) end) end + + local file_list = {} local embeddings = {} @@ -37,13 +58,25 @@ end) for i = 1, #file_list do local file_name = file_list[i] - -- TODO read and output to temp file, skipping YAML - -- local file_contents = utility.open(file_name, "r", function(file) - -- return file:read("*all") - -- end) + -- 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 tmp_file_name = utility.tmp_file_name() + utility.open(tmp_file_name, "w", function(file) + file:write(tmp_file_contents) + end) local embedding_model = "nomic-embed-text" - local output = utility.capture_safe("cat " .. file_name:enquote() .. " | ollama run " .. embedding_model) + 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 output = setmetatable({ vector = output }, { @@ -57,7 +90,7 @@ for i = 1, #file_list do print("Finished " .. i .. "/" .. #file_list .. " (" .. math.floor(i / #file_list * 100) .. "%)") end -local file = utility.open("dump.txt", "w", function(file) +utility.open("embeddings.json", "w", function(file) local output = json.encode(embeddings, { indent = true }) file:write(output) file:write("\n") diff --git a/test.lua b/test.lua index 6ee6420..a44e1ab 100755 --- a/test.lua +++ b/test.lua @@ -17,24 +17,6 @@ end -local tree -tree = function(path, fn) - utility.list(path or ".", function(path_name) - if utility.is_file(path_name) then - fn(path_name) - else - tree(path .. utility.path_separator .. path_name, fn) - end - end) -end - -tree(".", function(path_name) - print(path_name) -end) -os.exit() - - - local prompt = [[Return JSON: category, tags (array), summary (short)]] -- local prompt = [[Return YAML: category, tags (array), summary (short)]] @@ -50,28 +32,31 @@ prompt = prompt .. file_contents local model = "gemma3:4b" -- local output = utility.capture_safe("ollama run " .. model .. " --nowordwrap " .. prompt:enquote()) -local embedding_model = "nomic-embed-text" -local output = utility.capture_safe("ollama run " .. embedding_model .. " " .. file_contents:enquote()) +-- local embedding_model = "nomic-embed-text" +-- local output = utility.capture_safe("ollama run " .. embedding_model .. " " .. file_contents:enquote()) -output = output:sub(1, -2) -- strip extra newline from utility.capture_safe +-- output = output:sub(1, -2) -- strip extra newline from utility.capture_safe -print(output) - --- print(strip_markdown(output)) - -local embedding_metatable = { - __tojson = function(self, state) - return self.vector +-- 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 output = { - file_name = file_name, - vector = setmetatable({ vector = output }, embedding_metatable), -} +-- print(output) -print(json.encode(output, { indent = true })) - --- local file = utility.open("dump.txt", "w", function(file) --- file:write(output) --- end) +print(strip_frontmatter(file_contents)) +-- print(strip_markdown(output))