diff --git a/src/ml/imf/model.ts b/src/ml/imf/model.ts index 6a3bff2..389d13f 100644 --- a/src/ml/imf/model.ts +++ b/src/ml/imf/model.ts @@ -11,7 +11,7 @@ import type { Tensor } from "../types.js" import { verifyAndRead, parseManifest, type IMFManifest } from "./loader.js" import { resolve } from "./registry.js" import { normalizeArabicInput, repetitionGuardCut } from "./guards.js" -import { EOS_ID, PAD_ID, decode, encode } from "./tokens.js" +import { EOS_ID, PAD_ID, decode, encode, firstAlternates } from "./tokens.js" import { translateWindowed } from "./windows.js" interface InputMeta { @@ -102,7 +102,7 @@ export class IMFModel { const tokens = this.kv ? await this.greedyKv(hidden, maxLen, opts) : await this.greedyPlain(hidden, maxLen, opts) - return decode(tokens) + return firstAlternates(decode(tokens)) } /** Normalized text -> encoder hidden states; null when the input diff --git a/src/ml/imf/tokens.ts b/src/ml/imf/tokens.ts index 7705583..3deb463 100644 --- a/src/ml/imf/tokens.ts +++ b/src/ml/imf/tokens.ts @@ -30,3 +30,14 @@ export function decode(tokenIds: readonly number[]): string { } return decoder.decode(new Uint8Array(out)) } + +/** Seq2seq output contract: models trained on multi-reference corpora + * emit '/'-separated alternates per token; the runtime returns the + * primary choice only. */ +export function firstAlternates(text: string): string { + return text + .split(/\s+/) + .filter((t) => t.length > 0) + .map((t) => t.split("/")[0]) + .join(" ") +} diff --git a/test/ml/first-alternates.test.ts b/test/ml/first-alternates.test.ts new file mode 100644 index 0000000..f730982 --- /dev/null +++ b/test/ml/first-alternates.test.ts @@ -0,0 +1,18 @@ +import { describe, expect, it } from "vitest" + +import { firstAlternates } from "../../src/ml/imf/tokens.js" + +describe("firstAlternates", () => { + it("keeps the primary choice per token", () => { + expect(firstAlternates("وِكَالَةُ/وِكَالَةُ اَلْأَرْصَادِ")).toBe("وِكَالَةُ اَلْأَرْصَادِ") + }) + + it("leaves plain text untouched", () => { + expect(firstAlternates("مرحبا بالعالم")).toBe("مرحبا بالعالم") + }) + + it("handles edges", () => { + expect(firstAlternates("")).toBe("") + expect(firstAlternates("a/b c/d")).toBe("a c") + }) +})