@effect-uai/retrieval 0.12.1

package/src/HuggingFaceTokenizer.ts

@@ -0,0 +1,119 @@
+/**
+ * The core `Tokenizer` tag over `@huggingface/tokenizers`, which reads a
+ * model's `tokenizer.json` and needs no native build. The package is an
+ * optional peer dependency: `pnpm add @huggingface/tokenizers`.
+ *
+ * Fetching and construction are separate on purpose. `download` gives you
+ * plain JSON to keep wherever your application keeps things; `fromDefinition`
+ * builds the tokenizer from what you kept. Nothing here touches the filesystem
+ * or caches behind your back. `layer` composes the two for scripts happy to
+ * download on every run.
+ */
+import { Array as Arr, Effect, Layer, Redacted, Schema } from "effect"
+import { HttpClient, HttpClientRequest } from "effect/unstable/http"
+import { Tokenizer, type TokenizerService } from "@effect-uai/core/Tokenizer"
+import { Tokenizer as HfTokenizer } from "@huggingface/tokenizers"
+
+/** The vocabulary could not be fetched, parsed, or loaded. */
+export class TokenizerLoadError extends Schema.TaggedError<TokenizerLoadError>(
+  "@effect-uai/retrieval/TokenizerLoadError",
+)("TokenizerLoadError", {
+  model: Schema.optional(Schema.String),
+  reason: Schema.String,
+  cause: Schema.optional(Schema.Unknown),
+}) {}
+
+/**
+ * Everything that defines a tokenizer's behaviour, as published. Plain JSON:
+ * write it to disk, a row, or your bundle, and decode it back with this schema.
+ */
+export const Definition = Schema.Struct({
+  /** Contents of `tokenizer.json`. */
+  tokenizer: Schema.Unknown,
+  /** Contents of `tokenizer_config.json`, absent for repos that publish none. */
+  config: Schema.optional(Schema.Unknown),
+})
+
+export type Definition = typeof Definition.Type
+
+export type Options = {
+  /** A Hugging Face repo id, e.g. `"jinaai/jina-embeddings-v3"` or `"Xenova/gpt-4o"`. */
+  readonly model: string
+  /** Branch, tag, or commit. Default `"main"`. */
+  readonly revision?: string
+  /**
+   * A Hugging Face access token. Required for gated repos, where the Hub
+   * serves the files only to accounts that have accepted the model's terms
+   * (Google's tokenizers among them), and for private ones.
+   */
+  readonly token?: Redacted.Redacted
+}
+
+const fetchJson = (
+  options: Options,
+  file: string,
+): Effect.Effect<unknown, TokenizerLoadError, HttpClient.HttpClient> =>
+  Effect.gen(function* () {
+    const client = yield* HttpClient.HttpClient
+    const url = `https://huggingface.co/${options.model}/resolve/${options.revision ?? "main"}/${file}`
+    const failed = (reason: string) => (cause: unknown) =>
+      new TokenizerLoadError({ model: options.model, reason, cause })
+
+    const request = HttpClientRequest.get(url).pipe(
+      options.token === undefined
+        ? (self) => self
+        : HttpClientRequest.bearerToken(Redacted.value(options.token)),
+    )
+    const response = yield* client
+      .execute(request)
+      .pipe(Effect.mapError(failed(`could not fetch ${file}`)))
+    // 401/403 on a gated repo means no token, or terms not yet accepted.
+    return response.status >= 400
+      ? yield* new TokenizerLoadError({
+          model: options.model,
+          reason: `${url} returned ${response.status}`,
+        })
+      : yield* response.json.pipe(Effect.mapError(failed(`${file} is not valid JSON`)))
+  })
+
+/**
+ * Fetch a model's definition from the Hugging Face Hub. Cache the result
+ * yourself; a repo without a `tokenizer_config.json` yields `config: undefined`.
+ */
+export const download = (
+  options: Options,
+): Effect.Effect<Definition, TokenizerLoadError, HttpClient.HttpClient> =>
+  Effect.all({
+    tokenizer: fetchJson(options, "tokenizer.json"),
+    config: fetchJson(options, "tokenizer_config.json").pipe(Effect.orElseSucceed(() => undefined)),
+  })
+
+/** Build the tokenizer from a definition you already hold. No IO. */
+export const fromDefinition = (
+  definition: Definition,
+): Layer.Layer<Tokenizer, TokenizerLoadError> =>
+  Layer.effect(
+    Tokenizer,
+    Effect.try({
+      try: (): TokenizerService => {
+        const hf = new HfTokenizer(
+          definition.tokenizer as object,
+          (definition.config ?? {}) as object,
+        )
+        return {
+          encode: (text) => hf.encode(text).ids,
+          // The library rejects an empty array, but `encode("")` returns one and
+          // `decode` inverts `encode`. Its signature also wants a mutable array,
+          // though it mutates nothing.
+          decode: (tokens) => (tokens.length === 0 ? "" : hf.decode(Arr.copy(tokens))),
+        }
+      },
+      catch: (cause) => new TokenizerLoadError({ reason: "could not load the tokenizer", cause }),
+    }),
+  )
+
+/** Download and build in one step. For scripts; production should cache. */
+export const layer = (
+  options: Options,
+): Layer.Layer<Tokenizer, TokenizerLoadError, HttpClient.HttpClient> =>
+  Layer.unwrap(Effect.map(download(options), fromDefinition))

package/src/Rank.test.ts

@@ -0,0 +1,68 @@
+import { describe, expect, expectTypeOf, it } from "vitest"
+import { type Fused, rrf } from "./Rank.js"
+
+// 1-based rank r in a ranking of weight w contributes w / (k + r).
+const contribution = (rank: number, k = 60, weight = 1) => weight / (k + rank)
+
+describe("rrf", () => {
+  it("sums contributions across rankings and sorts descending", () => {
+    // "b" places 2nd and 1st; "a" places 1st and 3rd; "c" only 2nd.
+    const fused = rrf([
+      ["a", "b"],
+      ["b", "c", "a"],
+    ])
+    expect(fused.map((f) => f.value)).toEqual(["b", "a", "c"])
+    expect(fused[0]?.score).toBeCloseTo(contribution(2) + contribution(1), 12)
+    expect(fused[1]?.score).toBeCloseTo(contribution(1) + contribution(3), 12)
+    expect(fused[2]?.score).toBeCloseTo(contribution(2), 12)
+  })
+
+  it("scores an item present in only one ranking from that ranking alone", () => {
+    const fused = rrf([["a"], ["b"]])
+    expect(fused.map((f) => f.score)).toEqual([contribution(1), contribution(1)])
+  })
+
+  it("k damps the head of each list", () => {
+    // At k=1 a rank-1 hit is worth 1/2 and a rank-3 hit 1/4, so the single
+    // top placement beats the item that places 2nd twice. At k=60 it loses.
+    const rankings = [
+      ["a", "b"],
+      ["c", "b"],
+    ]
+    expect(rrf(rankings, { k: 1 })[0]?.value).toBe("b")
+    expect(rrf(rankings, { k: 1 })[0]?.score).toBeCloseTo(2 * contribution(2, 1), 12)
+    expect(rrf(rankings, { k: 0.5 })[0]?.value).toBe("b")
+  })
+
+  it("weights scale each ranking's contribution", () => {
+    const fused = rrf([["a"], ["b"]], { weights: [3, 1] })
+    expect(fused.map((f) => f.value)).toEqual(["a", "b"])
+    expect(fused[0]?.score).toBeCloseTo(contribution(1, 60, 3), 12)
+    // A missing weight falls back to 1 rather than dropping the ranking.
+    expect(rrf([["a"], ["b"]], { weights: [3] })[1]?.score).toBeCloseTo(contribution(1), 12)
+  })
+
+  it("breaks score ties by first-seen order", () => {
+    // Both place 1st in one ranking, so scores are identical.
+    const fused = rrf([["b"], ["a"]])
+    expect(fused[0]?.score).toBe(fused[1]?.score)
+    expect(fused.map((f) => f.value)).toEqual(["b", "a"])
+  })
+
+  it("returns empty for no rankings and for empty rankings", () => {
+    expect(rrf([])).toEqual([])
+    expect(rrf([[], []])).toEqual([])
+  })
+
+  it("fuses non-string items by Map identity", () => {
+    const fused = rrf([
+      [1, 2],
+      [2, 1],
+    ])
+    expectTypeOf(fused).toEqualTypeOf<Array<Fused<number>>>()
+    expect(fused.map((f) => f.score)).toEqual([
+      contribution(1) + contribution(2),
+      contribution(2) + contribution(1),
+    ])
+  })
+})

package/src/Rank.ts

@@ -0,0 +1,38 @@
+/**
+ * Rank fusion: merge several ranked lists of the same items into one
+ * ranking by position, so retrievers on incomparable score scales
+ * (BM25, cosine) can be combined without normalizing them.
+ */
+import { Array as Arr, Order, pipe } from "effect"
+
+export type Fused<A> = {
+  readonly value: A
+  readonly score: number
+}
+
+const byScore = Order.mapInput(Order.flip(Order.Number), (f: { readonly score: number }) => f.score)
+
+/**
+ * Reciprocal rank fusion: `score(v) = Σ weightᵢ / (k + rankᵢ(v))` over
+ * 1-based ranks, where a list omitting `v` contributes nothing. `k`
+ * defaults to 60, `weights` to 1 per ranking.
+ *
+ * Items are keyed by `Map` semantics, so fuse ids rather than fresh
+ * objects. Sorted descending; ties keep first-seen order.
+ */
+export const rrf = <A>(
+  rankings: ReadonlyArray<ReadonlyArray<A>>,
+  options?: { readonly k?: number; readonly weights?: ReadonlyArray<number> },
+): Array<Fused<A>> => {
+  const k = options?.k ?? 60
+  return pipe(
+    Arr.reduce(rankings, new Map<A, number>(), (scores, ranking, i) =>
+      Arr.reduce(ranking, scores, (acc, value, rank) =>
+        acc.set(value, (acc.get(value) ?? 0) + (options?.weights?.[i] ?? 1) / (k + rank + 1)),
+      ),
+    ),
+    Arr.fromIterable,
+    Arr.map(([value, score]): Fused<A> => ({ value, score })),
+    Arr.sort(byScore),
+  )
+}

package/src/index.ts

@@ -0,0 +1,2 @@
+export * as Chunking from "./Chunking.js"
+export * as Rank from "./Rank.js"