Skip to content

Feature/chunked embeddings #203

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
donhardman merged 1 commit into from
Sep 2, 2026
Merged

Feature/chunked embeddings #203

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 27 additions & 8 deletions knn/embeddings.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,7 +195,7 @@ class TextToEmbeddings_c : public TextToEmbeddings_i
TextToEmbeddings_c ( const ModelSettings_t & tSettings ) : m_tSettings ( tSettings ) {}

bool Initialize ( std::shared_ptr<LoadedLib_c> pLib, std::string & sError );
bool Convert ( const std::vector<std::string_view> & dTexts, std::vector<std::vector<float>> & dEmbeddings, std::string & sError, int iThreads = 0 ) const override;
bool Convert ( const std::vector<std::string_view> & dTexts, std::vector<std::vector<float>> & dEmbeddings, std::string & sError, int iThreads = 0, const knn::ChunkSettings_t * pChunk = nullptr, std::vector<size_t> * pRowOffsets = nullptr ) const override;
int GetDims() const override;

private:
Expand Down Expand Up @@ -255,7 +255,7 @@ bool TextToEmbeddings_c::Initialize ( std::shared_ptr<LoadedLib_c> pLib, std::st
}


bool TextToEmbeddings_c::Convert ( const std::vector<std::string_view> & dTexts, std::vector<std::vector<float>> & dEmbeddings, std::string & sError, int iThreads ) const
bool TextToEmbeddings_c::Convert ( const std::vector<std::string_view> & dTexts, std::vector<std::vector<float>> & dEmbeddings, std::string & sError, int iThreads, const knn::ChunkSettings_t * pChunk, std::vector<size_t> * pRowOffsets ) const
{
std::vector<StringItem> dStringItems;
for ( const auto & i : dTexts )
Expand All @@ -264,20 +264,25 @@ bool TextToEmbeddings_c::Convert ( const std::vector<std::string_view> & dTexts,
auto * pFuncs = m_pLib->GetLibFuncs();
assert(pFuncs);

ChunkSettings tChunk;
if ( pChunk )
{
tChunk.strategy = (uint32_t)pChunk->m_eStrategy;
tChunk.max_tokens = pChunk->m_uMaxTokens;
tChunk.overlap_tokens = pChunk->m_uOverlapTokens;
tChunk.max_chunks = pChunk->m_uMaxChunks;
}

// iThreads: 0 = use all available CPUs (default), >0 = cap worker count in the embeddings lib
// nullptr ChunkSettings = truncate strategy (today's behavior). Pass a populated
// ChunkSettings (e.g. STRATEGY_MEAN) once the table DDL exposes a chunking option.
FloatVecResult tVecResult = pFuncs->make_vect_embeddings ( &m_pModel, dStringItems.data(), dStringItems.size(), nullptr, iThreads );
// nullptr ChunkSettings = truncate strategy, i.e. one vector per input text
FloatVecResult tVecResult = pFuncs->make_vect_embeddings ( &m_pModel, dStringItems.data(), dStringItems.size(), pChunk ? &tChunk : nullptr, iThreads );
if ( tVecResult.m_szError )
{
sError = tVecResult.m_szError;
pFuncs->free_vec_result(tVecResult);
return false;
}

// truncate/mean return one vector per input row, so tVecResult.len == rows
// and this 1:1 copy is correct. A future multi-vector strategy (N vectors
// per row) would instead group via tVecResult.m_pRowOffsets[i..i+1].
dEmbeddings.resize ( tVecResult.len );
for ( size_t i = 0; i < tVecResult.len; i++ )
{
Expand All @@ -286,6 +291,20 @@ bool TextToEmbeddings_c::Convert ( const std::vector<std::string_view> & dTexts,
memcpy ( dEmbeddings[i].data(), tVec.ptr, sizeof(float)*tVec.len );
}

if ( pRowOffsets )
{
if ( !tVecResult.m_pRowOffsets || tVecResult.rows!=dTexts.size() )
{
sError = util::FormatStr ( "embeddings library returned a malformed result: %lld rows for %lld input texts%s", (long long)tVecResult.rows, (long long)dTexts.size(), tVecResult.m_pRowOffsets ? "" : ", no row offsets" );
pFuncs->free_vec_result(tVecResult);
return false;
}

pRowOffsets->resize ( tVecResult.rows+1 );
for ( size_t i = 0; i <= tVecResult.rows; i++ )
(*pRowOffsets)[i] = (size_t)tVecResult.m_pRowOffsets[i];
}

pFuncs->free_vec_result(tVecResult);

return true;
Expand Down
31 changes: 29 additions & 2 deletions knn/knn.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
namespace knn
{

static const int LIB_VERSION = 17;
static const int LIB_VERSION = 18;
static const uint32_t STORAGE_VERSION = 4;

enum class HNSWSimilarity_e
Expand Down Expand Up @@ -65,6 +65,32 @@ struct ModelSettings_t
bool m_bUseGPU = false;
};

/// how a document's text becomes one or many vectors. TRUNCATE/MEAN yield exactly one vector per
/// document and work with single-vector storage; FIXED/RECURSIVE/SENTENCE yield N and need an
/// attribute that can hold several vectors per row
enum class ChunkStrategy_e
{
TRUNCATE = 0,
MEAN = 1,
FIXED = 2,
RECURSIVE = 3,
SENTENCE = 4
};

inline bool IsMultiVectorStrategy ( ChunkStrategy_e eStrategy )
{
return eStrategy==ChunkStrategy_e::FIXED || eStrategy==ChunkStrategy_e::RECURSIVE || eStrategy==ChunkStrategy_e::SENTENCE;
}

/// mirrors the embeddings lib's ChunkSettings
struct ChunkSettings_t
{
ChunkStrategy_e m_eStrategy = ChunkStrategy_e::TRUNCATE;
uint32_t m_uMaxTokens = 0; // 0 = the model's own max input length
uint32_t m_uOverlapTokens = 0; // 0 = no overlap
uint32_t m_uMaxChunks = 0; // 0 = unlimited
};

struct AttrWithSettings_t : public common::SchemaAttr_t, public IndexSettings_t {};
using Schema_t = std::vector<AttrWithSettings_t>;

Expand Down Expand Up @@ -149,7 +175,8 @@ class TextToEmbeddings_i
public:
virtual ~TextToEmbeddings_i() = default;

virtual bool Convert ( const std::vector<std::string_view> & dTexts, std::vector<std::vector<float>> & dEmbeddings, std::string & sError, int iThreads = 0 ) const = 0;
// text i owns dEmbeddings[(*pRowOffsets)[i] .. (*pRowOffsets)[i+1]] in multi-vector case
virtual bool Convert ( const std::vector<std::string_view> & dTexts, std::vector<std::vector<float>> & dEmbeddings, std::string & sError, int iThreads = 0, const ChunkSettings_t * pChunk = nullptr, std::vector<size_t> * pRowOffsets = nullptr ) const = 0;
virtual int GetDims() const = 0;
};

Expand Down