Spaces:
Running
Running
File size: 4,512 Bytes
ca97aa9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
import { pipeline, FeatureExtractionPipeline } from "../../src/transformers.js";
import { MAX_MODEL_LOAD_TIME, MAX_TEST_EXECUTION_TIME, MAX_MODEL_DISPOSE_TIME, DEFAULT_MODEL_OPTIONS } from "../init.js";
const PIPELINE_ID = "feature-extraction";
export default () => {
describe("Feature Extraction", () => {
const model_id = "hf-internal-testing/tiny-random-BertModel";
const texts = ["This is a simple test.", "Hello world"];
/** @type {FeatureExtractionPipeline} */
let pipe;
beforeAll(async () => {
pipe = await pipeline(PIPELINE_ID, model_id, DEFAULT_MODEL_OPTIONS);
}, MAX_MODEL_LOAD_TIME);
it("should be an instance of FeatureExtractionPipeline ", () => {
expect(pipe).toBeInstanceOf(FeatureExtractionPipeline);
});
describe("batch_size=1", () => {
it(
"default",
async () => {
const output = await pipe(texts[0]);
expect(output.dims).toEqual([1, 20, 32]);
expect(output.type).toEqual("float32");
expect(output.mean().item()).toBeCloseTo(-1.538501215314625e-9, 6);
},
MAX_TEST_EXECUTION_TIME,
);
it(
"w/ cls pooling",
async () => {
const output = await pipe(texts[0], { pooling: "cls" });
expect(output.dims).toEqual([1, 32]);
expect(output.type).toEqual("float32");
expect(output.mean().item()).toBeCloseTo(2.491287887096405e-8, 6);
},
MAX_TEST_EXECUTION_TIME,
);
it(
"w/ mean pooling & normalization",
async () => {
const output = await pipe(texts[0], { pooling: "mean", normalize: true });
expect(output.dims).toEqual([1, 32]);
expect(output.type).toEqual("float32");
expect(output.mean().item()).toBeCloseTo(-2.0245352061465383e-9, 6);
},
MAX_TEST_EXECUTION_TIME,
);
it(
"w/ mean pooling & binary quantization",
async () => {
const output = await pipe(texts[0], { pooling: "mean", quantize: true, precision: "binary" });
expect(output.dims).toEqual([1, 32 / 8]);
expect(output.type).toEqual("int8");
expect(output.mean().item()).toEqual(-15);
},
MAX_TEST_EXECUTION_TIME,
);
it("w/ cls pooling & ubinary quantization", async () => {
const output = await pipe(texts[0], { pooling: "cls", quantize: true, precision: "ubinary" });
expect(output.dims).toEqual([1, 32 / 8]);
expect(output.type).toEqual("uint8");
expect(output.mean().item()).toEqual(140);
});
});
describe("batch_size>1", () => {
it(
"default",
async () => {
const output = await pipe(texts);
expect(output.dims).toEqual([texts.length, 20, 32]);
expect(output.type).toEqual("float32");
expect(output.mean().item()).toBeCloseTo(2.345950544935249e-9, 6);
},
MAX_TEST_EXECUTION_TIME,
);
it(
"w/ cls pooling",
async () => {
const output = await pipe(texts, { pooling: "cls" });
expect(output.dims).toEqual([texts.length, 32]);
expect(output.type).toEqual("float32");
expect(output.mean().item()).toBeCloseTo(1.6298145055770874e-8, 6);
},
MAX_TEST_EXECUTION_TIME,
);
it(
"w/ mean pooling & normalization",
async () => {
const output = await pipe(texts, { pooling: "mean", normalize: true });
expect(output.dims).toEqual([texts.length, 32]);
expect(output.type).toEqual("float32");
expect(output.mean().item()).toBeCloseTo(-1.538609240014921e-10, 6);
},
MAX_TEST_EXECUTION_TIME,
);
it("w/ mean pooling & binary quantization", async () => {
const output = await pipe(texts, { pooling: "mean", quantize: true, precision: "binary" });
expect(output.dims).toEqual([texts.length, 32 / 8]);
expect(output.type).toEqual("int8");
expect(output.mean().item()).toEqual(-14);
});
it("w/ cls pooling & ubinary quantization", async () => {
const output = await pipe(texts, { pooling: "cls", quantize: true, precision: "ubinary" });
expect(output.dims).toEqual([texts.length, 32 / 8]);
expect(output.type).toEqual("uint8");
expect(output.mean().item()).toEqual(140);
});
});
afterAll(async () => {
await pipe.dispose();
}, MAX_MODEL_DISPOSE_TIME);
});
};
|