rixty
/
SillyTavern
mirror of https://github.com/SillyTavern/SillyTavern.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Use Express router for classify endpoint

This commit is contained in:
valadaptive 2023-12-04 12:45:17 -05:00
parent 414c9bd5fb
commit 0ad753f317
2 changed files with 43 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
server.js
View File

@ -3618,7 +3618,7 @@ require('./src/endpoints/vectors').registerEndpoints(app, jsonParser);
require('./src/endpoints/translate').registerEndpoints(app, jsonParser);
// Emotion classification
require('./src/endpoints/classify').registerEndpoints(app, jsonParser);
app.use('/api/extra/classify', require('./src/endpoints/classify').router);
// Image captioning
require('./src/endpoints/caption').registerEndpoints(app, jsonParser);

87
src/endpoints/classify.js
View File

@ -1,53 +1,50 @@
const express = require('express');
const { jsonParser } = require('../express-common');
const TASK = 'text-classification';
/**
* @param {import("express").Express} app
* @param {any} jsonParser
*/
function registerEndpoints(app, jsonParser) {
const cacheObject = {};
const router = express.Router();
app.post('/api/extra/classify/labels', jsonParser, async (req, res) => {
try {
const module = await import('../transformers.mjs');
const pipe = await module.default.getPipeline(TASK);
const result = Object.keys(pipe.model.config.label2id);
return res.json({ labels: result });
} catch (error) {
console.error(error);
return res.sendStatus(500);
}
});
const cacheObject = {};
app.post('/api/extra/classify', jsonParser, async (req, res) => {
try {
const { text } = req.body;
router.post('/labels', jsonParser, async (req, res) => {
try {
const module = await import('../transformers.mjs');
const pipe = await module.default.getPipeline(TASK);
const result = Object.keys(pipe.model.config.label2id);
return res.json({ labels: result });
} catch (error) {
console.error(error);
return res.sendStatus(500);
}
});
async function getResult(text) {
if (Object.hasOwn(cacheObject, text)) {
return cacheObject[text];
} else {
const module = await import('../transformers.mjs');
const pipe = await module.default.getPipeline(TASK);
const result = await pipe(text, { topk: 5 });
result.sort((a, b) => b.score - a.score);
cacheObject[text] = result;
return result;
}
router.post('/', jsonParser, async (req, res) => {
try {
const { text } = req.body;
async function getResult(text) {
if (Object.hasOwn(cacheObject, text)) {
return cacheObject[text];
} else {
const module = await import('../transformers.mjs');
const pipe = await module.default.getPipeline(TASK);
const result = await pipe(text, { topk: 5 });
result.sort((a, b) => b.score - a.score);
cacheObject[text] = result;
return result;
}
console.log('Classify input:', text);
const result = await getResult(text);
console.log('Classify output:', result);
return res.json({ classification: result });
} catch (error) {
console.error(error);
return res.sendStatus(500);
}
});
}
module.exports = {
registerEndpoints,
};
console.log('Classify input:', text);
const result = await getResult(text);
console.log('Classify output:', result);
return res.json({ classification: result });
} catch (error) {
console.error(error);
return res.sendStatus(500);
}
});
module.exports = { router };