63 lines
1.7 KiB
JavaScript
63 lines
1.7 KiB
JavaScript
|
|
import http from 'http';
|
|
import querystring from 'querystring';
|
|
import url from 'url';
|
|
|
|
import { pipeline, env } from '@xenova/transformers';
|
|
|
|
class MyClassificationPipeline {
|
|
static task = 'text-classification';
|
|
static model = 'Xenova/distilbert-base-uncased-finetuned-sst-2-english';
|
|
static instance = null;
|
|
|
|
static async getInstance(progress_callback = null) {
|
|
if (this.instance === null) {
|
|
|
|
// NOTE: Uncomment this to change the cache directory
|
|
// env.cacheDir = './.cache';
|
|
|
|
this.instance = pipeline(this.task, this.model, { progress_callback });
|
|
}
|
|
|
|
return this.instance;
|
|
}
|
|
}
|
|
|
|
// Comment out this line if you don't want to start loading the model as soon as the server starts.
|
|
// If commented out, the model will be loaded when the first request is received (i.e,. lazily).
|
|
MyClassificationPipeline.getInstance();
|
|
|
|
// Define the HTTP server
|
|
const server = http.createServer();
|
|
const hostname = '127.0.0.1';
|
|
const port = 3000;
|
|
|
|
// Listen for requests made to the server
|
|
server.on('request', async (req, res) => {
|
|
// Parse the request URL
|
|
const parsedUrl = url.parse(req.url);
|
|
|
|
// Extract the query parameters
|
|
const { text } = querystring.parse(parsedUrl.query);
|
|
|
|
// Set the response headers
|
|
res.setHeader('Content-Type', 'application/json');
|
|
|
|
let response;
|
|
if (parsedUrl.pathname === '/classify' && text) {
|
|
const classifier = await MyClassificationPipeline.getInstance();
|
|
response = await classifier(text);
|
|
res.statusCode = 200;
|
|
} else {
|
|
response = { 'error': 'Bad request' }
|
|
res.statusCode = 400;
|
|
}
|
|
|
|
// Send the JSON response
|
|
res.end(JSON.stringify(response));
|
|
});
|
|
|
|
server.listen(port, hostname, () => {
|
|
console.log(`Server running at http://${hostname}:${port}/`);
|
|
});
|