{"version":3,"sources":["../../src/index.ts"],"sourcesContent":["import { validateConfigSchema } from \"@chroma-core/ai-embeddings-common\";\nimport { pipeline, ProgressCallback } from \"@huggingface/transformers\";\nimport { env as TransformersEnv } from \"@huggingface/transformers\";\n\nexport type DType =\n  | \"auto\"\n  | \"fp32\"\n  | \"fp16\"\n  | \"q8\"\n  | \"int8\"\n  | \"uint8\"\n  | \"q4\"\n  | \"bnb4\"\n  | \"q4f16\";\n\nexport type Quantization = DType | Record<string, DType>;\n\nexport type EmbeddingFunctionSpace = \"cosine\" | \"l2\" | \"ip\";\n\nexport interface DefaultEmbeddingFunctionConfig {\n  model_name?: string;\n  revision?: string;\n  dtype?: Quantization;\n  quantized?: boolean;\n  wasm?: boolean;\n}\n\nexport interface DefaultEmbeddingFunctionArgs {\n  modelName?: string;\n  revision?: string;\n  dtype?: Quantization;\n  /** @deprecated Use 'dtype' instead. If set to true, dtype value will be 'uint8' */\n  quantized?: boolean;\n  wasm?: boolean;\n}\n\nexport class DefaultEmbeddingFunction {\n  public readonly name: string = \"default\";\n  private readonly modelName: string;\n  private readonly revision: string;\n  private readonly dtype: Quantization | undefined;\n  private readonly quantized: boolean;\n  private readonly progressCallback: ProgressCallback | undefined = undefined;\n  private readonly wasm: boolean;\n\n  constructor(\n    args: Partial<\n      DefaultEmbeddingFunctionArgs & {\n        progressCallback: ProgressCallback | undefined;\n      }\n    > = {},\n  ) {\n    const {\n      modelName = \"Xenova/all-MiniLM-L6-v2\",\n      revision = \"main\",\n      dtype = undefined,\n      progressCallback = undefined,\n      quantized = false,\n      wasm = false,\n    } = args;\n\n    this.modelName = modelName;\n    this.revision = revision;\n    this.dtype = dtype || (quantized ? \"uint8\" : \"fp32\");\n    this.quantized = quantized;\n    this.progressCallback = progressCallback;\n    this.wasm = wasm;\n    if (this.wasm) {\n      TransformersEnv.backends.onnx.backend = \"wasm\";\n    }\n  }\n\n  public static buildFromConfig(\n    config: DefaultEmbeddingFunctionConfig,\n  ): DefaultEmbeddingFunction {\n    return new DefaultEmbeddingFunction(config);\n  }\n\n  public async generate(texts: string[]): Promise<number[][]> {\n    const pipe = await pipeline(\"feature-extraction\", this.modelName, {\n      revision: this.revision,\n      progress_callback: this.progressCallback,\n      dtype: this.dtype,\n    });\n\n    const output = await pipe(texts, { pooling: \"mean\", normalize: true });\n    return output.tolist();\n  }\n\n  public defaultSpace(): EmbeddingFunctionSpace {\n    return \"cosine\";\n  }\n\n  public supportedSpaces(): EmbeddingFunctionSpace[] {\n    return [\"cosine\", \"l2\", \"ip\"];\n  }\n\n  public getConfig(): DefaultEmbeddingFunctionConfig {\n    return {\n      model_name: this.modelName,\n      revision: this.revision,\n      dtype: this.dtype,\n      quantized: this.quantized,\n    };\n  }\n\n  public validateConfigUpdate(newConfig: DefaultEmbeddingFunctionConfig): void {\n    if (this.getConfig().model_name !== newConfig.model_name) {\n      throw new Error(\n        \"The DefaultEmbeddingFunction's 'model' cannot be changed after initialization.\",\n      );\n    }\n  }\n\n  public static validateConfig(config: DefaultEmbeddingFunctionConfig): void {\n    validateConfigSchema(config, \"transformers\");\n  }\n}\n"],"mappings":";;;;;;;;;;;;;;;;;;;;AAAA;AAAA;AAAA;AAAA;AAAA;AAAA,kCAAqC;AACrC,0BAA2C;AAC3C,IAAAA,uBAAuC;AAkChC,IAAM,2BAAN,MAAM,0BAAyB;AAAA,EASpC,YACE,OAII,CAAC,GACL;AAdF,SAAgB,OAAe;AAK/B,SAAiB,mBAAiD;AAUhE,UAAM;AAAA,MACJ,YAAY;AAAA,MACZ,WAAW;AAAA,MACX,QAAQ;AAAA,MACR,mBAAmB;AAAA,MACnB,YAAY;AAAA,MACZ,OAAO;AAAA,IACT,IAAI;AAEJ,SAAK,YAAY;AACjB,SAAK,WAAW;AAChB,SAAK,QAAQ,UAAU,YAAY,UAAU;AAC7C,SAAK,YAAY;AACjB,SAAK,mBAAmB;AACxB,SAAK,OAAO;AACZ,QAAI,KAAK,MAAM;AACb,2BAAAC,IAAgB,SAAS,KAAK,UAAU;AAAA,IAC1C;AAAA,EACF;AAAA,EAEA,OAAc,gBACZ,QAC0B;AAC1B,WAAO,IAAI,0BAAyB,MAAM;AAAA,EAC5C;AAAA,EAEA,MAAa,SAAS,OAAsC;AAC1D,UAAM,OAAO,UAAM,8BAAS,sBAAsB,KAAK,WAAW;AAAA,MAChE,UAAU,KAAK;AAAA,MACf,mBAAmB,KAAK;AAAA,MACxB,OAAO,KAAK;AAAA,IACd,CAAC;AAED,UAAM,SAAS,MAAM,KAAK,OAAO,EAAE,SAAS,QAAQ,WAAW,KAAK,CAAC;AACrE,WAAO,OAAO,OAAO;AAAA,EACvB;AAAA,EAEO,eAAuC;AAC5C,WAAO;AAAA,EACT;AAAA,EAEO,kBAA4C;AACjD,WAAO,CAAC,UAAU,MAAM,IAAI;AAAA,EAC9B;AAAA,EAEO,YAA4C;AACjD,WAAO;AAAA,MACL,YAAY,KAAK;AAAA,MACjB,UAAU,KAAK;AAAA,MACf,OAAO,KAAK;AAAA,MACZ,WAAW,KAAK;AAAA,IAClB;AAAA,EACF;AAAA,EAEO,qBAAqB,WAAiD;AAC3E,QAAI,KAAK,UAAU,EAAE,eAAe,UAAU,YAAY;AACxD,YAAM,IAAI;AAAA,QACR;AAAA,MACF;AAAA,IACF;AAAA,EACF;AAAA,EAEA,OAAc,eAAe,QAA8C;AACzE,0DAAqB,QAAQ,cAAc;AAAA,EAC7C;AACF;","names":["import_transformers","TransformersEnv"]}