rixty/SillyTavern
1
0
Fork 0
mirror of https://github.com/SillyTavern/SillyTavern.git synced 2025-06-05 21:59:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1454 from valadaptive/classify-router

Browse Source 
Use Express router for classify endpoint
This commit is contained in:
Cohee
2023-12-04 21:42:40 +02:00
committed by GitHub
parent c0cf274e6f a2387130f3
commit 88396935cd
2 changed files with 43 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
server.js
View File

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

87
src/endpoints/classify.js
View File

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