docs: custom embedding function for ts (#1479)

This commit is contained in:
Cory Grinstead
2024-07-30 18:19:55 -05:00
committed by GitHub
parent 277b753fd8
commit a062a92f6b
4 changed files with 295 additions and 170 deletions

View File

@@ -21,6 +21,7 @@ import {
Float32,
FloatLike,
type IntoVector,
Utf8,
isDataType,
isFixedSizeList,
isFloat,
@@ -192,6 +193,38 @@ export abstract class EmbeddingFunction<
}
}
/**
* an abstract class for implementing embedding functions that take text as input
*/
export abstract class TextEmbeddingFunction<
M extends FunctionOptions = FunctionOptions,
> extends EmbeddingFunction<string, M> {
//** Generate the embeddings for the given texts */
abstract generateEmbeddings(
texts: string[],
// biome-ignore lint/suspicious/noExplicitAny: we don't know what the implementor will do
...args: any[]
): Promise<number[][] | Float32Array[] | Float64Array[]>;
async computeQueryEmbeddings(data: string): Promise<Awaited<IntoVector>> {
return this.generateEmbeddings([data]).then((data) => data[0]);
}
embeddingDataType(): FloatLike {
return new Float32();
}
override sourceField(): [DataTypeLike, Map<string, EmbeddingFunction>] {
return super.sourceField(new Utf8());
}
computeSourceEmbeddings(
data: string[],
): Promise<number[][] | Float32Array[] | Float64Array[]> {
return this.generateEmbeddings(data);
}
}
export interface FieldOptions<T extends DataType = DataType> {
datatype: T;
dims?: number;