Merge pull request 'Refresh sources' (#2) from refresh_sources into main

Reviewed-on: #2

.gitignore

@@ -1,3 +1,3 @@
 .DS_Store
-PRIVATE_DATA/**
-!PRIVATE_DATA/.gitkeep
+PRIVATE_DATA/
+config.json

ReadMe.md

@@ -4,20 +4,60 @@ Finding useful things to do with local LLMs.
 All scripts work with data within `PRIVATE_DATA/` so that private data can't be
 accidentally committed.
 
-###### `cosine_similarity.lua`
+### `cosine_similarity.lua`
 Opens `embeddings.json` and creates `similarities.json` with a sorted list of
 comparisons.
 
-###### `generate_embeddings.lua`
+### `generate_embeddings.lua`
 Opens every file on its whitelist within the `notebook` directory, and generates
 embeddings (placed in `embeddings.json`). When files are too long, it truncates
 them.
 
-###### `least_similar.lua`
+### `least_similar.lua`
 Opens `similarities.json`, reverses the sort order, and saves it as
 `differences.json`.
 
-###### `synopsis_generator.lua`
+### `refresh_sources.lua`
+Creates/Maintains a store of chunked data with embeddings based on configurable
+sources. Example:
+
+```json
+{
+  "source name":{
+    "filters":{
+      "blacklist":[".git"],
+      "extension_whitelist":["md"]
+    },
+    "initialize_command":"git clone REMOTE .",
+    "path":"will be created before initialize_command is run",
+    "strip_frontmatter":true,
+    "refresh_command":"git reset --hard origin/main"
+  }
+}
+```
+
+The filters are based on `utility.tree`'s filter options (optional).
+`initialize_command` is only run the first time (optional),
+while `refresh_command` is run each time (optional).
+Commands will be run in the specified `path`.
+`strip_frontmatter` will remove YAML frontmatter (common in Markdown files).
+
+The embeddings are stored like so:
+
+```json
+{
+  "files":{
+    "PRIVATE_DATA/source_path/path/to/file.ext":["sha512sum", "another sum"]
+  },
+  "vectors":{
+    "sha512sum":[0.5, 0, 1, -0.5, -1, ...]
+  }
+}
+```
+
+TODO
+
+### `synopsis_generator.lua`
 Chooses a random file within `notebook`, and generates a novel synopsis from it.
 
 Arguments:

lib/log.lua

@@ -0,0 +1,75 @@
+local stored_messages = {}
+local display_levels = {}
+local enable_logging = true
+local message_limit = math.huge
+
+-- USAGE:
+--   -- set arbitrary log levels to be printed immediately
+--   log{ info = true, warning = true, ducksauce = true, }
+--   -- send anything (except nil) to a log level
+--   log("bacom", true, "text", 5, function() end, {})
+--   -- print all messages saved at a particular log level
+--   log("error")
+--   -- toggle ALL logging on and off with a boolean
+--   log(false) -- this deletes all stored messages
+--   -- set a maximum number of stored messages per log level
+--   log(50) -- warning: reaching large limits makes logging expensive
+--   -- do anything to the stored messages (this example deletes them all)
+--  log(function(messages) return {} end)
+
+return function(options, ...)
+  local options_type = type(options)
+
+  if (options_type == "string") and enable_logging then
+    if not stored_messages[options] then stored_messages[options] = {} end
+    local message_table = stored_messages[options]
+
+    -- turn log message into text
+    local tab = {...}
+    for i = 1, #tab do
+      tab[i] = tostring(tab[i])
+    end
+    local current_message = table.concat(tab, "\t")
+
+    if #current_message == 0 then
+      -- print all stored_messages of specified level
+      for i = 1, #message_table do
+        print(message_table[i])
+      end
+
+    else
+      -- store message (and print if in display_levels)
+      message_table[#message_table + 1] = current_message
+      if #message_table > message_limit then table.remove(message_table, 1) end
+      if display_levels[options] then print(current_message) end
+    end
+
+  -- set display_levels
+  elseif options_type == "table" then
+    for k,v in pairs(options) do
+      display_levels[k] = v
+    end
+
+  elseif options_type == "number" then
+    message_limit = options
+    for k,v in pairs(stored_messages) do
+      while #v > message_limit do
+        table.remove(v, 1)
+      end
+    end
+
+  -- turn all logging off and delete stored messages; or turn it on
+  elseif options_type == "boolean" then
+    if options then
+      enable_logging = true
+    else
+      enable_logging = false
+      stored_messages = {}
+    end
+
+  -- arbitrary access :D
+  elseif options_type == "function" then
+    local result = options(stored_messages)
+    if type(result) == "table" then stored_messages = result end
+  end
+end

lib/timing.lua

@@ -1,18 +1,26 @@
 local timing = {}
 
+local function human_readable_time(delta)
+  if delta >= 2*7*24*60*60 then -- if more than 2 weeks
+    delta = tostring(math.floor(delta/(7*24*60*60))/10) .. " weeks"
+  elseif delta >= 2*24*60*60 then -- if more than 2 days
+    delta = tostring(math.floor(delta/(24*60*60))/10) .. " days"
+  elseif delta >= 2*60*60 then -- if more than 2 hours
+    delta = tostring(math.floor(delta/(60*60/10))/10) .. " hours"
+  elseif delta >= 2*60 then -- if more then 2 minutes
+    delta = tostring(math.floor(delta/(60/10))/10) .. " minutes"
+  else
+    delta = tostring(delta) .. " seconds"
+  end
+  return delta
+end
+
 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
+    local delta = human_readable_time(time.time - previous.time)
 
     print(delta, previous.label)
   end
@@ -32,6 +40,13 @@ timing.mark = function(label)
   if #timing > 1 then
     timing.display(#timing)
   end
+  print("", "", label)
+end
+
+timing.estimate = function(current_position, total_operations)
+  local delta = os.time() - timing[#timing].time
+  local estimate = delta * total_operations / current_position - delta
+  return human_readable_time(estimate)
 end
 
 return timing

lib/utility.lua

@@ -18,7 +18,7 @@ if package.config:sub(1, 1) == "\\" then
   }
 else
   utility = {
-    OS = "UNIX-like",
+    OS = "Linux",
     path_separator = "/",
     temp_directory = "/tmp/",
     commands = {
@@ -81,6 +81,10 @@ standard_library_addition(string, "split", function(s, delimiter)
   return result
 end)
 
+utility.leftpad = function(text, length, character)
+  return string.rep(character or " ", length - #(tostring(text))) .. text
+end
+
 
 
 utility.require = function(...)
@@ -217,7 +221,7 @@ utility.list = function(path, func)
 
   local run = function(fn)
     for line in output:gmatch("[^\r\n]+") do -- thanks to https://stackoverflow.com/a/32847589
-      if not (line == "." or line == "..") then
+      if not ((line == ".") or (line == "..")) then
         fn(line)
       end
     end
@@ -235,7 +239,8 @@ utility.ls = function(...)
   return utility.list(...)
 end
 
-utility.tree = function(path, options, fn)
+local tree
+tree = function(path, options, fn)
   if type(options) == "function" then
     fn = options
     options = {}
@@ -245,16 +250,25 @@ utility.tree = function(path, options, fn)
     if options.blacklist and options.blacklist[path_name] then return end
     if options.whitelist and (not options.whitelist[path_name]) then return end
 
-    if options.extension_blacklist or options.extension_whitelist then
-      local _, _, extension = utility.split_path_components(path_name)
-      if options.extension_blacklist and options.extension_blacklist[extension] then return end
-      if options.extension_whitelist and (not options.extension_whitelist[extension]) then return end
-    end
-
     if utility.is_file(path_name) then
+      if options.extension_blacklist or options.extension_whitelist then
+        local _, _, extension = utility.split_path_components(path_name)
+        if options.extension_blacklist and options.extension_blacklist[extension] then return end
+        if options.extension_whitelist and (not options.extension_whitelist[extension]) then return end
+      end
+
       fn(path_name)
     else
-      utility.tree(path .. utility.path_separator .. path_name, options, fn)
+      tree(path .. utility.path_separator .. path_name, options, fn)
+    end
+  end)
+end
+utility.tree = function(path, options, fn)
+  tree(path, options, function(path_name)
+    if path_name:find(path) == 1 then
+      fn(path_name)
+    else
+      fn(path .. utility.path_separator .. path_name)
     end
   end)
 end
@@ -265,10 +279,11 @@ utility.read_file = function(file_name)
   end)
 end
 
-utility.write_file = function(file_name, text)
+utility.write_file = function(file_name, ...)
+  local text = table.concat{...}
   return utility.open(file_name, "w", function(file)
     file:write(text)
-    file:write("\n")
+    -- file:write("\n") -- I need to make sure /I/ handle this instead of trying to automate it
   end)
 end
 
@@ -299,6 +314,16 @@ utility.file_size = function(file_path)
   return utility.open(file_path, "rb", function(file) return file:seek("end") end)
 end
 
+utility.sha512sum = function(file_path)
+  local sha512sum
+  if (utility.OS == "Linux") or (utility.OS == "macOS") then
+    sha512sum = utility.capture_safe("shasum -U -a 512 " .. file_path:enquote())
+  elseif utility.OS == "Windows" then
+    error("utility.sha512sum() not implemented for Windows.")
+  end
+  return sha512sum:sub(1, 128)
+end
+
 
 
 utility.escape_quotes_and_escapes = function(input)
@@ -570,4 +595,8 @@ end
 
 
 
+if (utility.OS == "Linux") and (utility.capture_safe("uname"):find("Darwin") == 1) then
+  utility.OS = "macOS"
+end
+
 return utility

refresh_sources.lua

@@ -0,0 +1,262 @@
+#!/usr/bin/env luajit
+
+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 log = utility.require("log")
+local text_processing = utility.require("text_processing")
+local timing = utility.require("timing")
+
+-- TODO make utility have a function for getting/setting defaults where locking is only used to set defaults if they aren't present
+-- TODO there should be a check function for a lock so a warning/error can be dumped?
+local config = utility.get_config("no-lock")
+if not config.models then
+  config.models = {
+    embedding = {
+      model = "qwen3-embedding:0.6b",
+      max_chunk_size = 32768,
+    },
+    initialized_sources = {},
+  }
+  utility.save_config()
+end
+
+log{
+  info = true,
+  warning = true,
+  -- debug = true,
+  -- files = true, -- debugging why the wrong files are selected
+  -- sha = true, -- what the fuck is going on with sha sums?
+}
+
+local memory_path = "PRIVATE_DATA" .. utility.path_separator .. "memory"
+local embeddings_file_path = memory_path .. utility.path_separator .. "+embeddings.json"
+local tmp_file_path = "PRIVATE_DATA" .. utility.path_separator .. ".tmp.2b65c19b-0883-49ca-8247-b1fe7760f922"
+
+local embeddings
+if not utility.path_exists(embeddings_file_path) then
+  os.execute("mkdir -p " .. memory_path:enquote())
+  utility.save_data({
+    files = {},
+    vectors = {},
+  }, embeddings_file_path)
+end
+embeddings = utility.load_data(embeddings_file_path)
+local function embeddings_debug()
+  log("debug", "Embeddings loaded.", embeddings, embeddings.files, embeddings.vectors)
+  local file_count = 0
+  for k,v in pairs(embeddings.files) do
+    file_count = file_count + 1
+  end
+  local vector_count = 0
+  for k,v in pairs(embeddings.vectors) do
+    vector_count = vector_count + 1
+  end
+  log("debug", file_count .. " files.")
+  log("debug", vector_count .. " vectors.")
+  -- os.exit(1)
+end
+embeddings_debug()
+
+
+
+local refresh_file_list = function(source_name, data_source)
+  timing.mark("Assembling file list for " .. source_name .. ".")
+
+  local full_path = "PRIVATE_DATA" .. utility.path_separator .. data_source.path
+  if data_source.initialize_command and (not config.models.initialized_sources[data_source.path]) then
+    -- NOTE this is where we'd want to check/obtain a lock on the config so we can safely run this
+    os.execute("mkdir -p " .. full_path:enquote() .. " && cd " .. full_path:enquote() .. " && " .. data_source.initialize_command)
+    config.models.initialized_sources[data_source.path] = true
+    utility.save_config()
+  end
+  if data_source.refresh_command then
+    os.execute("cd " .. full_path:enquote() .. " && " .. data_source.refresh_command)
+  end
+
+  local compiled_filters = {}
+  if data_source.filters then
+    for name, object in pairs(data_source.filters) do
+      compiled_filters[name] = utility.enumerate(object)
+    end
+  end
+
+  local file_list = {}
+  utility.tree("PRIVATE_DATA" .. utility.path_separator .. data_source.path, compiled_filters, function(file_name)
+    log("files", file_name)
+    file_list[#file_list + 1] = file_name
+  end)
+
+  timing.mark("Finished assembling file list for " .. source_name .. ".")
+  return file_list
+end
+
+-- returns nothing when too much text is sent
+local generate_embeddings = function(text)
+  if #text > config.models.embedding.max_chunk_size then
+    return nil, "generate_embeddings() must only be passed appropriately-sized chunks!"
+  end
+
+  local result = utility.llm_prompt(text, config.models.embedding.model)
+  return json.decode(result)
+end
+
+-- returns nothing for empty files and errors
+local process_file = function(data_source, file_name)
+  local text = utility.read_file(file_name)
+
+  if data_source.strip_frontmatter then
+    text = text_processing.strip_frontmatter(text)
+  end
+
+  if #text == 0 then
+    log("empty", file_name .. "\n is empty and being skipped.")
+    return
+  end
+
+  local chunk_size = config.models.embedding.max_chunk_size
+  local half_chunk_size = math.floor(chunk_size / 2)
+  local chunks = { text }
+
+  while #text > chunk_size do
+    local first_chunk = text:sub(1, chunk_size)
+    local overlap_chunk = text:sub(half_chunk_size, chunk_size + half_chunk_size - 1)
+
+    chunks[#chunks + 1] = first_chunk
+    chunks[#chunks + 1] = overlap_chunk
+
+    text = text:sub(chunk_size)
+    if (#text > half_chunk_size) and (not (#text > chunk_size)) then
+      -- last chunk would be skipped if we didn't handle this here
+      chunks[#chunks + 1] = text
+    end
+  end
+
+  local new_embeddings = {}
+  for i = 1, #chunks do
+    log("debug", "Embedding length:", #chunks[i])
+    new_embeddings[i] = generate_embeddings(chunks[i]) or {}
+  end
+
+  if #new_embeddings[1] == 0 then
+    log("debug", "Vector lengths:")
+    for e = 1, #new_embeddings do
+      log("debug", "", e, #new_embeddings[e])
+    end
+    if #new_embeddings == 1 then
+      -- Ollama very rarely errors with:
+      --   Error: do embedding request: Post "http://127.0.0.1:53441/v1/embeddings": EOF
+      -- but it is inconsistent and re-running will eventually fix it.
+      log("warning", file_name .. "\n encountered an embedding error and will be skipped this run only.")
+      return
+    end
+
+    -- average all embeddings to make the core file embedding
+    local count = #new_embeddings[2]
+    for vector_index = 1, count do
+      local total = 0
+      for chunk = 2, #new_embeddings do
+        total = total + new_embeddings[chunk][vector_index]
+      end
+      new_embeddings[1][vector_index] = total / count
+    end
+  end
+
+  return chunks, new_embeddings
+end
+
+-- returns nothing for errors
+local memorize_file = function(data_source, file_name)
+  local file_chunks, file_embeddings = process_file(data_source, file_name)
+  if not file_chunks then return end
+
+  local file_sums = {}
+  for i = 1, #file_chunks do
+    local function loop()
+      local text = file_chunks[i]
+      local current_embedding = file_embeddings[i]
+
+      utility.write_file(tmp_file_path, text)
+
+      local sha512sum = utility.sha512sum(tmp_file_path)
+      log("sha", "New sum? " .. sha512sum, tostring(embeddings.vectors[sha512sum]))
+      file_sums[#file_sums + 1] = sha512sum
+      if embeddings.vectors[sha512sum] then
+        log("sha", "Detected previously extant sum, skipping.")
+        return
+      end
+
+      os.execute(utility.commands.move .. tmp_file_path:enquote()
+        .. " " .. (memory_path .. utility.path_separator .. sha512sum):enquote())
+      embeddings.vectors[sha512sum] = current_embedding
+      log("sha", "Saved new sum!")
+    end
+    loop()
+  end
+
+  return file_sums
+end
+
+local refresh_sources = function()
+  os.execute("mkdir -p " .. memory_path:enquote())
+
+  timing.mark("Checking memory for missing embeddings.")
+  utility.list(memory_path, function(file_name)
+    log("files", file_name)
+    if (file_name == ".DS_Store") or (file_name == "+embeddings.json") then return end
+
+    local file_path = memory_path .. utility.path_separator .. file_name
+    local sha512sum = utility.sha512sum(file_path)
+    log("sha", "New sum? " .. sha512sum, tostring(embeddings.vectors[sha512sum]))
+    if not embeddings.vectors[sha512sum] then
+      log("sha", "Saved new sum!")
+      memorize_file({}, file_path)
+    end
+  end)
+  timing.mark("Finished checking memory for missing embeddings.")
+
+  local sources = utility.load_data("PRIVATE_DATA/sources.json")
+  for source_name, data_source in pairs(sources) do
+    local file_list = refresh_file_list(source_name, data_source)
+
+    timing.mark("Generating embeddings from source \"" .. source_name .. "\"")
+    for f = 1, #file_list do
+      local file_name = file_list[f]
+      local function loop()
+        local sha512sum = utility.sha512sum(file_name)
+        log("sha", file_name, sha512sum, "\n Sum present? " .. tostring(embeddings.vectors[sha512sum]))
+        if embeddings.vectors[sha512sum] then
+          log("debug", file_name .. "\n has already been embedded, skipping.")
+          -- add file reference if it was missing
+          if not embeddings.files[file_name] then
+            embeddings.files[file_name] = { sha512sum }
+          end
+          return
+        end
+
+        log("debug", "Generating embeddings for " .. file_name .. ".")
+        local file_sums = memorize_file(data_source, file_name)
+        embeddings.files[file_name] = file_sums
+      end
+      loop()
+      log("info", "Finished " .. utility.leftpad(f, #tostring(#file_list), "0")
+        .. "/" .. #file_list
+        .. " (" .. utility.leftpad(math.floor(f / #file_list * 100), 3, "0")
+        .. "%) ETA: " .. timing.estimate(f, #file_list))
+    end
+
+    timing.mark("Finished generating embeddings from source \"" .. source_name .. "\"")
+  end
+
+  utility.save_data(embeddings)
+  if utility.path_exists(tmp_file_path) then
+    os.execute("rm " .. tmp_file_path:enquote())
+  end
+end
+
+refresh_sources()
+timing.mark("Finished.")
+print("")
+timing.display()
+log("warning")

synopsis_generator.lua

@@ -108,7 +108,7 @@ local export_ordered_list_of_prompts = function()
     output[#output + 1] = "## Scoring\n\n```json\n" .. json.encode(item.scoring, { indent = true, }) .. "\n```\n"
   end
 
-  utility.write_file("PRIVATE_DATA/Ordered Synopses.md", table.concat(output, "\n"))
+  utility.write_file("PRIVATE_DATA/Ordered Synopses.md", table.concat(output, "\n"), "\n")
   os.execute("pandoc \"PRIVATE_DATA/Ordered Synopses.md\" -o \"PRIVATE_DATA/Ordered Synopses.epub\"")
 end
 

test.lua

@@ -6,16 +6,9 @@ local json = utility.require("dkjson")
 
 local file_name = arg[1]
 
+print(utility.OS)
 
-
-local function strip_markdown(text)
-  local tab = text:split("\n")
-  table.remove(tab, 1)    -- codeblock opening
-  table.remove(tab, #tab) -- end of codeblock
-  print(table.concat(tab, "\n"))
-end
-
-
+os.exit(0)
 
 local prompt = [[Return JSON: category, tags (array), summary (short)]]
 -- local prompt = [[Return YAML: category, tags (array), summary (short)]]
@@ -24,39 +17,6 @@ if prompt:sub(-1) ~= "\n" then
   prompt = prompt .. "\n\n"
 end
 
-local file_contents = utility.open(file_name, "r", function(file)
-  return file:read("*all")
-end)
+local file_contents = utility.read_file(file_name)
 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())
-
--- output = output:sub(1, -2) -- strip extra newline from utility.capture_safe
-
--- 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
-
--- print(output)
-
-print(strip_frontmatter(file_contents))
--- print(strip_markdown(output))
+-- print(utility.llm_prompt(prompt))