1
import express from 'express';
2

3
import { getPipeline } from '../transformers.js';
4

5
const TASK = 'text-classification';
6

7
export const router = express.Router();
8

9
/**
10
 * @type {Map<string, object>} Cache for classification results
11
 */
12
const cacheObject = new Map();
13

14
router.post('/labels', async (req, res) => {
15
    try {
16
        const pipe = await getPipeline(TASK);
17
        const result = Object.keys(pipe.model.config.label2id);
18
        return res.json({ labels: result });
19
    } catch (error) {
20
        console.error(error);
21
        return res.sendStatus(500);
22
    }
23
});
24

25
router.post('/', async (req, res) => {
26
    try {
27
        const { text } = req.body;
28

29
        /**
30
         * Get classification result for a given text
31
         * @param {string} text Text to classify
32
         * @returns {Promise<object>} Classification result
33
         */
34
        async function getResult(text) {
35
            if (cacheObject.has(text)) {
36
                return cacheObject.get(text);
37
            } else {
38
                const pipe = await getPipeline(TASK);
39
                const result = await pipe(text, { topk: 5 });
40
                result.sort((a, b) => b.score - a.score);
41
                cacheObject.set(text, result);
42
                return result;
43
            }
44
        }
45

46
        console.debug('Classify input:', text);
47
        const result = await getResult(text);
48
        console.debug('Classify output:', result);
49

50
        return res.json({ classification: result });
51
    } catch (error) {
52
        console.error(error);
53
        return res.sendStatus(500);
54
    }
55
});