diff --git a/bindings/node/lib/bindings/raw-encoding.test.ts b/bindings/node/lib/bindings/raw-encoding.test.ts index b0f2becd..580b66f5 100644 --- a/bindings/node/lib/bindings/raw-encoding.test.ts +++ b/bindings/node/lib/bindings/raw-encoding.test.ts @@ -1,5 +1,6 @@ import { promisify } from "util"; +import { PaddingDirection } from "./enums"; import { Model, WordPiece, WordPieceOptions } from "./models"; import { whitespacePreTokenizer } from "./pre-tokenizers"; import { RawEncoding } from "./raw-encoding"; @@ -137,4 +138,31 @@ describe("RawEncoding", () => { expect(index).toBeUndefined(); }); }); + + describe("pad", () => { + it("works correctly with only one parameter", () => { + encoding.pad(10); + expect(encoding.getTokens()).toHaveLength(10); + }); + + it("accepts `undefined` as second parameter", () => { + encoding.pad(10, undefined); + expect(encoding.getTokens()).toHaveLength(10); + }); + + it("accepts options as second parameter", () => { + encoding.pad(10, { + direction: PaddingDirection.Left, + padToken: "[PA]", + padTypeId: 10, + padId: 400 + }); + + const tokens = encoding.getTokens(); + expect(tokens).toHaveLength(10); + expect(tokens[0]).toBe("[PA]"); + expect(encoding.getTypeIds()[0]).toBe(10); + expect(encoding.getIds()[0]).toBe(400); + }); + }); }); diff --git a/bindings/node/native/src/encoding.rs b/bindings/node/native/src/encoding.rs index d77df772..a21a664b 100644 --- a/bindings/node/native/src/encoding.rs +++ b/bindings/node/native/src/encoding.rs @@ -257,7 +257,7 @@ declare_types! { // padToken?: string = "[PAD]" // } let length = cx.extract::(0)?; - let params = cx.extract_opt::(0)? + let params = cx.extract_opt::(1)? .map_or_else(tk::PaddingParams::default, |p| p.0); let mut this = cx.this();