VEMB RAW implemented.

This commit is contained in:
antirez 2025-02-05 11:21:06 +01:00
parent 8206c782b5
commit 6c1e55d07c
2 changed files with 42 additions and 10 deletions

View File

@ -138,6 +138,19 @@ Because vector sets perform insertion time normalization and optional
quantization, the returned vector could be approximated. `VEMB` will take
care to de-quantized and de-normalize the vector before returning it.
It is possible to ask VEMB to return raw data, that is, the interal representation used by the vector: fp32, int8, or a bitmap for binary quantization. This behavior is triggered by the `RAW` option of of VEMB:
VEMB word_embedding apple RAW
In this case the return value of the command is an array of three or more elements:
1. The name of the quantization used, that is one of: "fp32", "bin", "q8".
2. The a string blob containing the raw data, 4 bytes fp32 floats for fp32, a bitmap for binary quants, or int8 bytes array for q8 quants.
3. A float representing the l2 of the vector before normalization. You need to multiply by this vector if you want to de-normalize the value for any reason.
For q8 quantization, an additional elements is also returned: the quantization
range, so the integers from -127 to 127 represent (normalized) components
in the range `-range`, `+range`.
**VLINKS: introspection command that shows neighbors for a node**
VLINKS key element [WITHSCORES]

23
vset.c
View File

@ -860,8 +860,19 @@ int VREM_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
* reduced vector is returned instead. */
int VEMB_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
RedisModule_AutoMemory(ctx);
int raw_output = 0; // RAW option.
if (argc != 3) return RedisModule_WrongArity(ctx);
if (argc < 3) return RedisModule_WrongArity(ctx);
/* Parse arguments. */
for (int j = 3; j < argc; j++) {
const char *opt = RedisModule_StringPtrLen(argv[j], NULL);
if (!strcasecmp(opt,"raw")) {
raw_output = 1;
} else {
return RedisModule_ReplyWithError(ctx,"ERR invalid option");
}
}
/* Get key and element. */
RedisModuleString *key = argv[1];
@ -885,6 +896,14 @@ int VEMB_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
return RedisModule_ReplyWithNull(ctx);
}
if (raw_output) {
int output_qrange = vset->hnsw->quant_type == HNSW_QUANT_Q8;
RedisModule_ReplyWithArray(ctx, 3+output_qrange);
RedisModule_ReplyWithSimpleString(ctx, vectorSetGetQuantName(vset));
RedisModule_ReplyWithStringBuffer(ctx, node->vector, hnsw_quants_bytes(vset->hnsw));
RedisModule_ReplyWithDouble(ctx, node->l2);
if (output_qrange) RedisModule_ReplyWithDouble(ctx, node->quants_range);
} else {
/* Get the vector associated with the node. */
float *vec = RedisModule_Alloc(sizeof(float) * vset->hnsw->vector_dim);
hnsw_get_node_vector(vset->hnsw, node, vec); // May dequantize/denorm.
@ -893,8 +912,8 @@ int VEMB_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) {
RedisModule_ReplyWithArray(ctx, vset->hnsw->vector_dim);
for (uint32_t i = 0; i < vset->hnsw->vector_dim; i++)
RedisModule_ReplyWithDouble(ctx, vec[i]);
RedisModule_Free(vec);
}
return REDISMODULE_OK;
}