mtf-nlp/src/server.js

183 lines
4.0 KiB
JavaScript

// Import dotenv file if not in production
if (process.env.NODE_ENV !== "production") {
require("dotenv").config();
}
const {env, pipeline} = require("@xenova/transformers");
const path = require("path");
env.allowRemoteModels = false;
env.localModelPath = path.resolve(__dirname, "./models/transformers").replace(/\\/g, "/");
const models = require("./models");
const db = models.db;
const redis = models.redis;
const {sequelize, Sequelize} = db;
const Intent = db.intent;
async function getIntents() {
try {
const intents = await Intent.findAll();
console.log(intents);
return intents;
} catch (err) {
console.error("getIntents error:", err);
}
}
const config = require("./config");
/*const INTENTS = [
{
name: "play_music",
examples: ["play music", "start music", "resume music", "play songs"],
},
{
name: "stop_music",
examples: ["stop music", "pause music", "halt music", "stop playback"],
},
{
name: "set_volume",
examples: ["set volume to 30", "volume 50", "turn volume up", "turn volume down"],
},
];*/
let extractor;
function normalizeText(text) {
return text
.toLowerCase()
.replace(/[^\w\s]/g, " ")
.replace(/\s+/g, " ")
.trim();
}
function cosineSimilarity(a, b) {
let dot = 0;
let magA = 0;
let magB = 0;
for (let i = 0; i < a.length; i++) {
dot += a[i] * b[i];
magA += a[i] * a[i];
magB += b[i] * b[i];
}
return dot / (Math.sqrt(magA) * Math.sqrt(magB));
}
async function initModel() {
extractor = await pipeline("feature-extraction", `${config.NLP_TRANSFORMER}`, {
dtype: 'fp32',
quantized: false,
device: "cpu",
});
console.log("MiniLM model loaded");
}
async function embed(text) {
const output = await extractor(normalizeText(text), {
pooling: "mean",
normalize: true,
});
return Array.from(output.data);
}
async function classifyIntent(text) {
const inputEmbedding = await embed(text);
let best = {
name: "unknown",
confidence: 0,
matched: null,
};
for (const intent of INTENTS) {
for (const example of intent.examples) {
const exampleEmbedding = await embed(example);
const score = cosineSimilarity(inputEmbedding, exampleEmbedding);
if (score > best.confidence) {
best = {
name: intent.name,
confidence: score,
matched: example,
};
}
}
}
return best;
}
async function handleIntent(streamKey, msg) {
const text = msg.message?.text;
const deviceId = msg.message?.device_id;
if (!text) return;
const intent = await classifyIntent(text);
await redis.client.xAdd(streamKey, '*', {
transcript_id: msg.id,
command: intent.name,
confidence: String(intent.confidence),
matched: intent.matched || "",
device_id: deviceId || "",
text,
});
console.log("TODO: Handle Intents", intent);
}
async function processor(group, consumer, streamKey) {
await redis.client.connect().catch(console.error);
while (true) {
const res = await redis.client.xReadGroup(
group,
consumer,
[{key: streamKey, id: '>'}], // '>' means undelivered messages
{COUNT: 5, BLOCK:3000}
);
if (!res) continue;
for (const stream of res) {
for (const msg of stream.messages) {
try {
console.log(`[${consumer}] Processing:`, msg);
await handleIntent("commands", msg);
await redis.client.xAck(streamKey, group, msg.id);
} catch (err) {
console.error("Failed to process message:", err);
}
}
}
}
}
const main = async () => {
// Sync DB (not forced in production)
const doForce = config.FORCE_DB_SYNC === "true" && config.NODE_ENV !== "production";
await sequelize.authenticate();
await sequelize.sync({force: doForce});
if (doForce) console.info("Database dropped and resynced (force=true).");
const INTENTS = getIntents();
await initModel();
await redis.ensureGroup("transcripts", 'nlp');
await redis.ensureGroup("commands", 'nlp');
processor('nlp', `${config.DEVICE_ID}`, "transcripts").catch(console.error);
}
main().catch(console.error);