107 lines
3.1 KiB
TypeScript
107 lines
3.1 KiB
TypeScript
import Bun from 'bun';
|
|
import { spawn } from 'node:child_process';
|
|
import { parseArgs } from 'node:util';
|
|
import fs from 'fs';
|
|
import path from 'path';
|
|
import { factorize } from '../preprocess/factorize';
|
|
import { parseTja } from '../preprocess/parse'
|
|
|
|
const { values } = parseArgs({
|
|
args: Bun.argv,
|
|
options: {
|
|
workingDir: {
|
|
type: "string"
|
|
},
|
|
dataDir: {
|
|
type: "string"
|
|
},
|
|
script: {
|
|
type: "string"
|
|
},
|
|
trainSize: {
|
|
type: 'string'
|
|
},
|
|
validSize: {
|
|
type: 'string'
|
|
}
|
|
},
|
|
allowPositionals: true,
|
|
});
|
|
|
|
if (!values.dataDir || !values.workingDir || !values.script) {
|
|
console.error("Usage: bun run script/train_factor.ts --workingDir <dir> --dataDir <dir> --script <python_script> --trainSize <num> --validSize <num>");
|
|
process.exit(1);
|
|
}
|
|
|
|
// 1. factors.json 생성 (없을 경우)
|
|
generateFactors();
|
|
|
|
// 2. 파이썬 학습 스크립트 실행
|
|
console.log(`Starting training with factor-based script: ${values.script}`);
|
|
const child = spawn("python3", [
|
|
values.script,
|
|
"--workingDir", values.workingDir,
|
|
"--dataDir", values.dataDir,
|
|
"--trainSize", (Number(values.trainSize) || 1000).toString(),
|
|
"--validSize", (Number(values.validSize) || 200).toString(),
|
|
]);
|
|
|
|
child.stdout.pipe(process.stdout);
|
|
child.stderr.pipe(process.stderr);
|
|
|
|
child.on("close", (code) => {
|
|
console.log(`Training process exited with code ${code}`);
|
|
process.exit(code || 0);
|
|
});
|
|
|
|
// functions
|
|
function generateFactors() {
|
|
const workingDir = values.workingDir ?? '';
|
|
if (!fs.existsSync(workingDir)) fs.mkdirSync(workingDir, { recursive: true });
|
|
|
|
const factorsPath = path.join(workingDir, 'factors.json');
|
|
if (fs.existsSync(factorsPath)) {
|
|
console.log('factors.json already exists, skipping generation.');
|
|
return;
|
|
}
|
|
|
|
const dataDir = values.dataDir ?? '';
|
|
const tjaDir = path.join(dataDir, 'tja');
|
|
if (!fs.existsSync(tjaDir)) {
|
|
console.error(`TJA directory not found: ${tjaDir}`);
|
|
return;
|
|
}
|
|
|
|
const files = fs.readdirSync(tjaDir).filter(f => f.endsWith('.tja'));
|
|
console.log(`Generating factors from ${files.length} TJA files...`);
|
|
|
|
const results: any[] = [];
|
|
for (const file of files) {
|
|
try {
|
|
const tja = fs.readFileSync(path.join(tjaDir, file), 'utf-8');
|
|
const songno = path.basename(file, '.tja');
|
|
const parsed = parseTja(tja);
|
|
|
|
const courses = [
|
|
{ diff: 'oni', data: parsed?.oni },
|
|
{ diff: 'ura', data: parsed?.edit }
|
|
];
|
|
|
|
for (const course of courses) {
|
|
if (course.data) {
|
|
results.push({
|
|
songno,
|
|
difficulty: course.diff,
|
|
factors: factorize(course.data)
|
|
});
|
|
}
|
|
}
|
|
} catch (err) {
|
|
console.error(`Error processing ${file}:`, err);
|
|
}
|
|
}
|
|
|
|
fs.writeFileSync(factorsPath, JSON.stringify(results, null, 2), 'utf-8');
|
|
console.log(`factors.json generated at ${factorsPath}`);
|
|
}
|