example.lua

-- good resources
-- https://opensearch.org/blog/improving-document-retrieval-with-sparse-semantic-encoders/
-- https://huggingface.co/opensearch-project/opensearch-neural-sparse-encoding-v1
--
-- run with
-- text-embeddings-router --model-id opensearch-project/opensearch-neural-sparse-encoding-v1 --pooling splade

local cjson = require("cjson")
local http = require("socket.http")
local ltn12 = require("ltn12")
local pgmoon = require("pgmoon")
local pgvector = require("./src/pgvector")

local pg = pgmoon.new({
 database = "pgvector_example",
 user = os.getenv("USER")
})

assert(pg:connect())
assert(pg:query("CREATE EXTENSION IF NOT EXISTS vector"))
assert(pg:query("DROP TABLE IF EXISTS documents"))
assert(pg:query("CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding sparsevec(30522))"))

function embed(inputs)
 local url = "http://localhost:3000/embed_sparse"
 local data = {
 inputs = inputs,
 }
 local headers = {
 ["Content-Type"] = "application/json"
 }

 local chunks = {}
 local r, c, h = http.request {
 method = "POST",
 url = url,
 headers = headers,
 source = ltn12.source.string(cjson.encode(data)),
 sink = ltn12.sink.table(chunks)
 }
 assert(c == 200)
 local res = cjson.decode(table.concat(chunks))

 local embeddings = {}
 for i, item in ipairs(res) do
 local embedding = {}
 for i, v in ipairs(item) do
 embedding[v["index"] + 1] = v["value"]
 end
 embeddings[i] = embedding
 end
 return embeddings
end

local documents = {
 "The dog is barking",
 "The cat is purring",
 "The bear is growling"
}
local embeddings = embed(documents)
for i, content in ipairs(documents) do
 local embedding = embeddings[i]
 assert(pg:query("INSERT INTO documents (content, embedding) VALUES ($1, $2)", content, pgvector.sparsevec(embedding, 30522)))
end

local query = "forest"
local embedding = embed({query})[1]
local res = assert(pg:query("SELECT content FROM documents ORDER BY embedding <#> $1 LIMIT 5", pgvector.sparsevec(embedding, 30522)))
for i, row in ipairs(res) do
 print(row["content"])
end