scripts/generate-similarity-rankings-cache

#!/usr/bin/env luajit
-- Generate pre-sorted similarity rankings cache for fast HTML generation
-- Reads triangular individual similarity files and outputs a cache with
-- pre-sorted rankings for each poem (sorted by similarity, descending)
--
-- This is a post-processing step after GPU similarity generation.
-- The cache format mirrors the diversity cache for consistency.

-- {{{ setup_dir_path
local function setup_dir_path(provided_dir)
if provided_dir then
return provided_dir
end
return "/mnt/mtwo/programming/ai-stuff/neocities-modernization"
end
-- }}}

local DIR = setup_dir_path(arg and arg[1])

package.path = DIR .. "/libs/?.lua;" .. DIR .. "/src/?.lua;" .. package.path

local utils = require("utils")
local dkjson = require("dkjson")

utils.init_assets_root({DIR})

-- {{{ function load_all_triangular_similarities
-- Load all individual similarity files and reconstruct full symmetric matrix
local function load_all_triangular_similarities(model_name)
model_name = model_name or "embeddinggemma:latest"
local model_dir = model_name:gsub(":", "_")
-- Issue 10-054: similarities + ranking cache are movable (embeddings_dir, RAM).
local similarities_dir = utils.similarities_dir(model_name)

print(string.format("[SIMILARITY CACHE] Loading triangular files from: %s", similarities_dir))

-- Build full similarity map: poem_index -> {target_index -> similarity}
-- Using triangular property: similarity(A,B) = similarity(B,A)
local full_similarities = {}
local files_processed = 0
local start_time = os.time()

-- List similarity files - prefer poem_index_*.json (newer format)
-- Skip sorting for speed - order doesn't matter for cache building
local handle = io.popen(string.format("find '%s' -maxdepth 1 -name 'poem_index_*.json' 2>/dev/null | wc -l",
similarities_dir))
local poem_index_count = tonumber(handle:read("*a"):match("%d+")) or 0
handle:close()

local files = {}
if poem_index_count > 0 then
-- Use newer poem_index_*.json files (no sorting needed)
handle = io.popen(string.format("find '%s' -maxdepth 1 -name 'poem_index_*.json' 2>/dev/null",
similarities_dir))
for line in handle:lines() do
table.insert(files, line)
end
handle:close()
print(string.format("[SIMILARITY CACHE] Using poem_index_*.json files (%d found)", #files))
else
-- Fall back to legacy poem_.json files (excluding poem_index_)
handle = io.popen(string.format("find '%s' -maxdepth 1 -name 'poem_.json' ! -name 'poem_index_' 2>/dev/null",
similarities_dir))
for line in handle:lines() do
table.insert(files, line)
end
handle:close()
print(string.format("[SIMILARITY CACHE] Using legacy poem_*.json files (%d found)", #files))
end

local total_files = #files
print(string.format("[SIMILARITY CACHE] Found %d similarity files", total_files))

local skipped_count = 0
for _, filepath in ipairs(files) do
local data = utils.read_json_file(filepath)
if data and data.metadata and data.similarities then
local source_index = data.metadata.poem_index

-- Skip files without poem_index
if not source_index then
skipped_count = skipped_count + 1
goto continue
end

-- Initialize source poem's similarity map if needed
if not full_similarities[source_index] then
full_similarities[source_index] = {}
end

-- Process each similarity entry
for _, entry in ipairs(data.similarities) do
local target_index = tonumber(entry.id)
local similarity = entry.similarity

if target_index and similarity then
-- Store bidirectionally (symmetric matrix)
full_similarities[source_index][target_index] = similarity

-- Also store reverse direction
if not full_similarities[target_index] then
full_similarities[target_index] = {}
end
full_similarities[target_index][source_index] = similarity
end
end

files_processed = files_processed + 1

if files_processed % 500 == 0 then
local elapsed = os.time() - start_time
local rate = files_processed / math.max(elapsed, 1)
print(string.format("[SIMILARITY CACHE] Progress: %d/%d files (%.1f%%), %.1f files/sec",
files_processed, total_files,
(files_processed / total_files) * 100, rate))
end
end
::continue::
end

if skipped_count > 0 then
print(string.format("[SIMILARITY CACHE] Skipped %d files without poem_index", skipped_count))
end

print(string.format("[SIMILARITY CACHE] Loaded similarities for %d poems from %d files",
utils.table_count(full_similarities), files_processed))

return full_similarities
end
-- }}}

-- {{{ function utils.table_count
-- Count entries in a table (since # only works for arrays)
function utils.table_count(t)
local count = 0
for _ in pairs(t) do count = count + 1 end
return count
end
-- }}}

-- {{{ function generate_sorted_rankings
-- Convert full similarity map to pre-sorted rankings
local function generate_sorted_rankings(full_similarities)
print("[SIMILARITY CACHE] Generating sorted rankings...")

local rankings = {}
local total_poems = utils.table_count(full_similarities)
local processed = 0
local start_time = os.time()

for poem_index, similarities in pairs(full_similarities) do
-- Convert to array for sorting
local sim_array = {}
for target_index, similarity in pairs(similarities) do
table.insert(sim_array, {
id = target_index,
similarity = similarity
})
end

-- Sort by similarity descending (most similar first)
table.sort(sim_array, function(a, b)
return a.similarity > b.similarity
end)

-- Extract just the sorted IDs (similarity values not needed at runtime)
local sorted_ids = {}
for _, entry in ipairs(sim_array) do
table.insert(sorted_ids, entry.id)
end

rankings[tostring(poem_index)] = sorted_ids

processed = processed + 1
if processed % 1000 == 0 then
local elapsed = os.time() - start_time
print(string.format("[SIMILARITY CACHE] Sorted: %d/%d poems (%.1f%%)",
processed, total_poems, (processed / total_poems) * 100))
end
end

print(string.format("[SIMILARITY CACHE] Generated rankings for %d poems", processed))
return rankings
end
-- }}}

-- {{{ function save_rankings_cache
local function save_rankings_cache(rankings, model_name)
model_name = model_name or "embeddinggemma:latest"
local model_dir = model_name:gsub(":", "_")
-- Issue 10-054: similarity ranking cache is movable (embeddings_dir, RAM).
local cache_file = utils.embeddings_dir(model_name) .. "/similarity_rankings_cache.json"

local total_poems = utils.table_count(rankings)

local cache_data = {
metadata = {
total_poems = total_poems,
generated_at = os.date("%Y-%m-%d %H:%M:%S"),
algorithm = "gpu_vulkan_triangular_sorted",
format = "pre_sorted_rankings",
description = "Pre-sorted similarity rankings for fast HTML generation"
},
rankings = rankings
}

print(string.format("[SIMILARITY CACHE] Saving cache to: %s", cache_file))
local start_time = os.time()

utils.write_json_file(cache_file, cache_data)

local elapsed = os.time() - start_time
local file_handle = io.open(cache_file, "r")
local file_size = file_handle:seek("end")
file_handle:close()

print(string.format("[SIMILARITY CACHE] ✅ Cache saved: %.1f MB in %d seconds",
file_size / 1024 / 1024, elapsed))

return cache_file
end
-- }}}

-- {{{ main
local function main()
print("=" .. string.rep("=", 60))
print("[SIMILARITY CACHE] Pre-sorted Similarity Rankings Generator")
print("=" .. string.rep("=", 60))

local overall_start = os.time()

-- Step 1: Load all triangular similarity files
local full_similarities = load_all_triangular_similarities()

-- Step 2: Generate sorted rankings for each poem
local rankings = generate_sorted_rankings(full_similarities)

-- Step 3: Save to cache file
local cache_file = save_rankings_cache(rankings)

local total_time = os.time() - overall_start
print("")
print(string.format("[SIMILARITY CACHE] ✅ Complete! Total time: %.1f minutes", total_time / 60))
print(string.format("[SIMILARITY CACHE] Cache file: %s", cache_file))
end

main()
-- }}}