text-gen-bot/index.js

63 lines
2.2 KiB
JavaScript
Raw Normal View History

2022-08-15 17:28:08 -04:00
import config from './config.json' assert {type: "json"};
import { MatrixClient, SimpleFsStorageProvider, AutojoinRoomsMixin } from "matrix-bot-sdk";
2022-08-16 02:36:47 -04:00
import { spawn } from 'node:child_process';
2022-08-15 19:48:20 -04:00
import fs from "fs";
2022-08-15 17:28:08 -04:00
const storage = new SimpleFsStorageProvider("storage.json");
2022-08-16 00:47:38 -04:00
const client = new MatrixClient(config.homeserver, config.token, storage);
2022-08-17 16:23:49 -04:00
const pyFile = "textgen.py";
2022-08-15 17:28:08 -04:00
2022-08-17 16:23:49 -04:00
AutojoinRoomsMixin.setupOnClient(client);
2022-08-15 20:13:55 -04:00
client.start().then(() => console.log(`Client has started!\n`));
2022-08-15 17:28:08 -04:00
2022-08-15 19:48:20 -04:00
let messageCounter = 0;
2022-08-17 16:23:49 -04:00
let trainCounter = 0;
2022-08-15 19:48:20 -04:00
2022-08-15 17:28:08 -04:00
client.on("room.message", (roomId, event) => {
2022-08-16 00:47:38 -04:00
if (!event["content"] || event["sender"] === config.user) return;
2022-08-17 16:23:49 -04:00
2022-08-16 00:47:38 -04:00
++messageCounter;
2022-08-17 16:23:49 -04:00
++trainCounter;
let userMessage = event["content"]["body"].split(" ");
if (userMessage[0].startsWith(config.prefix)) {
userMessage[0] = userMessage[0].replace(config.prefix, '').toLowerCase();
} else {
2022-08-18 02:11:58 -04:00
fs.appendFile(config.file, userMessage.join(' ') + "\n", function (err) {
2022-08-17 16:23:49 -04:00
if (err) throw err;
});
};
// ? send message every N messages if big enough dataset is present
if ((!(messageCounter % config.frequency) && !(lineCount(config.file) < config.size)) || userMessage[0] === "speak") {
2022-08-16 00:47:38 -04:00
console.log("Generating message...");
2022-08-16 02:36:47 -04:00
2022-08-17 16:23:49 -04:00
userMessage.shift()
const python = spawn('python', [pyFile, "generate", userMessage.join(' ')]);
2022-08-16 02:36:47 -04:00
2022-08-17 16:23:49 -04:00
python.stdout.on('data', (message) => {
2022-08-16 02:36:47 -04:00
message = message.toString();
2022-08-17 16:23:49 -04:00
client.sendText(roomId, message); // ? send generated message to room
2022-08-16 02:36:47 -04:00
});
2022-08-17 16:23:49 -04:00
console.log("Message sent!");
2022-08-16 02:36:47 -04:00
python.on('close'); // ? close python process when finished
2022-08-15 19:48:20 -04:00
};
2022-08-16 03:45:11 -04:00
2022-08-17 16:23:49 -04:00
if (trainCounter >= config.retrain || userMessage[0] === "train") {
console.log("Retraining the AI...");
2022-08-16 03:45:11 -04:00
2022-08-17 16:23:49 -04:00
trainCounter = 0;
const python = spawn('python', [pyFile, "train"]);
python.stdout.on('data', function (message) {
console.log(message.toString());
});
console.log("Training finished!");
python.on('close'); // ? close python process when finished
};
2022-08-15 19:48:20 -04:00
});
2022-08-16 00:47:38 -04:00
function lineCount(text) {
return fs.readFileSync(text).toString().split("\n").length - 1;
2022-08-17 16:23:49 -04:00
};