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

Add ability to override local classification model

Browse Source
This commit is contained in:
Cohee
2023-09-11 01:25:22 +03:00
parent 5cc6a2dca6
commit c76c76410c
2 changed files with 55 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
src/classify.mjs
View File

@ -1,22 +1,37 @@
import { pipeline, TextClassificationPipeline, env } from 'sillytavern-transformers';
import { pipeline, env } from 'sillytavern-transformers';
import path from 'path';
// Limit the number of threads to 1 to avoid issues on Android
env.backends.onnx.wasm.numThreads = 1;
class PipelineAccessor {
/**
* @type {TextClassificationPipeline}
* @type {import("sillytavern-transformers").TextClassificationPipeline}
*/
pipe;
async get() {
if (!this.pipe) {
const cache_dir = path.join(process.cwd(), 'cache');
this.pipe = await pipeline('text-classification', 'Cohee/distilbert-base-uncased-go-emotions-onnx', { cache_dir, quantized: true });
const model = this.getClassificationModel();
this.pipe = await pipeline('text-classification', model, { cache_dir, quantized: true });
}
return this.pipe;
}
getClassificationModel() {
const DEFAULT_MODEL = 'Cohee/distilbert-base-uncased-go-emotions-onnx';
try {
const config = require(path.join(process.cwd(), './config.conf'));
const model = config?.extras?.classificationModel;
return model || DEFAULT_MODEL;
} catch {
console.warn('Failed to read config.conf, using default classification model.');
return DEFAULT_MODEL;
}
}
}
/**
@ -28,30 +43,40 @@ function registerEndpoints(app, jsonParser) {
const pipelineAccessor = new PipelineAccessor();
app.post('/api/extra/classify/labels', jsonParser, async (req, res) => {
const pipe = await pipelineAccessor.get();
const result = Object.keys(pipe.model.config.label2id);
return res.json({ labels: result });
try {
const pipe = await pipelineAccessor.get();
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) => {
const { text } = req.body;
try {
const { text } = req.body;
async function getResult(text) {
if (cacheObject.hasOwnProperty(text)) {
return cacheObject[text];
} else {
const pipe = await pipelineAccessor.get();
const result = await pipe(text);
cacheObject[text] = result;
return result;
async function getResult(text) {
if (cacheObject.hasOwnProperty(text)) {
return cacheObject[text];
} else {
const pipe = await pipelineAccessor.get();
const result = await pipe(text);
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);
}
console.log('Classify input:', text);
const result = await getResult(text);
console.log('Classify output:', result);
return res.json({ classification: result });
});
}