functional embeddings generation
This commit is contained in:
1 parent
2e4d3d9094
commit
de3b36cc93
1 file changed
+27
-2
+27
-2
@@ -7,6 +7,9 @@ local json = utility.require("dkjson")
|
|||||||
local PATH = "notebook"
|
local PATH = "notebook"
|
||||||
local whitelist = { md = true, }
|
local whitelist = { md = true, }
|
||||||
|
|
||||||
|
local embedding_model = "nomic-embed-text"
|
||||||
|
local maximum_file_size = 2048
|
||||||
|
|
||||||
-- strip YAML frontmatter (if present)
|
-- strip YAML frontmatter (if present)
|
||||||
-- can error, will return nil & error message
|
-- can error, will return nil & error message
|
||||||
local function strip_frontmatter(text)
|
local function strip_frontmatter(text)
|
||||||
@@ -37,6 +40,10 @@ tree = function(path, fn)
|
|||||||
end)
|
end)
|
||||||
end
|
end
|
||||||
|
|
||||||
|
local function leftpad(text, length)
|
||||||
|
return string.rep("0", length - #(tostring(text))) .. text
|
||||||
|
end
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
local file_list = {}
|
local file_list = {}
|
||||||
@@ -56,6 +63,7 @@ tree(PATH, function(file_name)
|
|||||||
end)
|
end)
|
||||||
|
|
||||||
for i = 1, #file_list do
|
for i = 1, #file_list do
|
||||||
|
local function _run()
|
||||||
local file_name = file_list[i]
|
local file_name = file_list[i]
|
||||||
|
|
||||||
-- stripping YAML frontmatter before generating embeddings
|
-- stripping YAML frontmatter before generating embeddings
|
||||||
@@ -69,16 +77,30 @@ for i = 1, #file_list do
|
|||||||
tmp_file_contents = file_contents
|
tmp_file_contents = file_contents
|
||||||
end
|
end
|
||||||
|
|
||||||
|
if #tmp_file_contents == 0 then
|
||||||
|
print(file_name .. "\n is empty and will be skipped.")
|
||||||
|
return
|
||||||
|
end
|
||||||
|
|
||||||
|
if #tmp_file_contents > maximum_file_size then
|
||||||
|
print(file_name .. "\n is too long. Only the first " .. maximum_file_size .. " bytes will be scanned.")
|
||||||
|
tmp_file_contents = tmp_file_contents:sub(1, maximum_file_size)
|
||||||
|
end
|
||||||
|
|
||||||
local tmp_file_name = utility.tmp_file_name()
|
local tmp_file_name = utility.tmp_file_name()
|
||||||
utility.open(tmp_file_name, "w", function(file)
|
utility.open(tmp_file_name, "w", function(file)
|
||||||
file:write(tmp_file_contents)
|
file:write(tmp_file_contents)
|
||||||
end)
|
end)
|
||||||
|
|
||||||
local embedding_model = "nomic-embed-text"
|
|
||||||
local output = utility.capture_safe("cat " .. tmp_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)
|
os.execute("rm " .. tmp_file_name)
|
||||||
output = output:sub(1, -2) -- strip extra newline from utility.capture_safe
|
output = output:sub(1, -2) -- strip extra newline from utility.capture_safe
|
||||||
|
|
||||||
|
if #output == 0 then
|
||||||
|
print("Warning: " .. file_name .. " did not have embeddings generated.")
|
||||||
|
return
|
||||||
|
end
|
||||||
|
|
||||||
output = setmetatable({ vector = output }, {
|
output = setmetatable({ vector = output }, {
|
||||||
__tojson = function(self, state)
|
__tojson = function(self, state)
|
||||||
return self.vector
|
return self.vector
|
||||||
@@ -87,7 +109,10 @@ for i = 1, #file_list do
|
|||||||
|
|
||||||
embeddings[file_name] = { vector = output }
|
embeddings[file_name] = { vector = output }
|
||||||
|
|
||||||
print("Finished " .. i .. "/" .. #file_list .. " (" .. math.floor(i / #file_list * 100) .. "%)")
|
print("Finished " .. leftpad(i, #tostring(#file_list)) .. "/" .. #file_list .. " (" .. leftpad(math.floor(i / #file_list * 100), 3) .. "%)")
|
||||||
|
end
|
||||||
|
|
||||||
|
_run()
|
||||||
end
|
end
|
||||||
|
|
||||||
utility.open("embeddings.json", "w", function(file)
|
utility.open("embeddings.json", "w", function(file)
|
||||||
|
|||||||
Reference in new issue
Block a user