Spaces:
Running
Running
File size: 3,387 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 |
import { pipeline, TextClassificationPipeline } 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 = "text-classification";
export default () => {
describe("Text Classification", () => {
const model_id = "hf-internal-testing/tiny-random-BertForSequenceClassification";
/** @type {TextClassificationPipeline} */
let pipe;
beforeAll(async () => {
pipe = await pipeline(PIPELINE_ID, model_id, DEFAULT_MODEL_OPTIONS);
}, MAX_MODEL_LOAD_TIME);
it("should be an instance of TextClassificationPipeline", () => {
expect(pipe).toBeInstanceOf(TextClassificationPipeline);
});
describe("batch_size=1", () => {
it(
"default (top_k=1)",
async () => {
const output = await pipe("a");
const target = [{ label: "LABEL_0", score: 0.5076976418495178 }];
expect(output).toBeCloseToNested(target, 5);
},
MAX_TEST_EXECUTION_TIME,
);
it(
"custom (top_k=2)",
async () => {
const output = await pipe("a", { top_k: 2 });
const target = [
{ label: "LABEL_0", score: 0.5076976418495178 },
{ label: "LABEL_1", score: 0.49230238795280457 },
];
expect(output).toBeCloseToNested(target, 5);
},
MAX_TEST_EXECUTION_TIME,
);
});
describe("batch_size>1", () => {
it(
"default (top_k=1)",
async () => {
const output = await pipe(["a", "b c"]);
const target = [
{ label: "LABEL_0", score: 0.5076976418495178 },
{ label: "LABEL_0", score: 0.5077522993087769 },
];
expect(output).toBeCloseToNested(target, 5);
},
MAX_TEST_EXECUTION_TIME,
);
it(
"custom (top_k=2)",
async () => {
const output = await pipe(["a", "b c"], { top_k: 2 });
const target = [
[
{ label: "LABEL_0", score: 0.5076976418495178 },
{ label: "LABEL_1", score: 0.49230238795280457 },
],
[
{ label: "LABEL_0", score: 0.5077522993087769 },
{ label: "LABEL_1", score: 0.49224773049354553 },
],
];
expect(output).toBeCloseToNested(target, 5);
},
MAX_TEST_EXECUTION_TIME,
);
it(
"multi_label_classification",
async () => {
const problem_type = pipe.model.config.problem_type;
pipe.model.config.problem_type = "multi_label_classification";
const output = await pipe(["a", "b c"], { top_k: 2 });
const target = [
[
{ label: "LABEL_0", score: 0.5001373887062073 },
{ label: "LABEL_1", score: 0.49243971705436707 },
],
[
{ label: "LABEL_0", score: 0.5001326203346252 },
{ label: "LABEL_1", score: 0.492380291223526 },
],
];
expect(output).toBeCloseToNested(target, 5);
// Reset problem type
pipe.model.config.problem_type = problem_type;
},
MAX_TEST_EXECUTION_TIME,
);
});
afterAll(async () => {
await pipe.dispose();
}, MAX_MODEL_DISPOSE_TIME);
});
};
|