DUB-AI: multi-task модель (тип озвучки + автор-спан + 18+), обучение/инференс/ONNX/сервис, тесты и документация
This commit is contained in:
DUB-AI committed
commit
8dbca3f4a3
2873 files changed
+44479
No files matched your search
@@ -0,0 +1,81 @@
|
||||
// Локальный инференс DUB-AI без Python (путь Б: ONNX), Node.js.
|
||||
//
|
||||
// Подготовка (один раз):
|
||||
// cd examples/typescript && npm init -y && npm i onnxruntime-node @huggingface/transformers
|
||||
// Экспорт модели: python -m src.export_onnx (создаст models/best-onnx)
|
||||
//
|
||||
// Запуск:
|
||||
// node onnx_local.mjs "../../models/best-onnx" "Перевод: Дублированный | MovieDalen"
|
||||
|
||||
import { readFileSync } from "node:fs";
|
||||
import ort from "onnxruntime-node";
|
||||
import { AutoTokenizer } from "@huggingface/transformers";
|
||||
|
||||
const [modelDir, textArg] = process.argv.slice(2);
|
||||
if (!modelDir || !textArg) {
|
||||
console.error('использование: node onnx_local.mjs <models/best-onnx> "текст дорожки"');
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
const labelMaps = JSON.parse(readFileSync(`${modelDir}/label_maps.json`, "utf-8"));
|
||||
const { aliases } = JSON.parse(readFileSync(`${modelDir}/aliases.json`, "utf-8"));
|
||||
|
||||
const tokenizer = await AutoTokenizer.from_pretrained(modelDir);
|
||||
const session = await ort.InferenceSession.create(`${modelDir}/model.onnx`);
|
||||
|
||||
// на входе здесь одна дорожка; сегментацию страниц делает predict.mjs (путь А)
|
||||
const inputs = tokenizer(textArg, { truncation: true, max_length: 128 });
|
||||
|
||||
const feeds = {};
|
||||
for (const name of ["input_ids", "attention_mask", "token_type_ids"]) {
|
||||
if (inputs[name] !== undefined) {
|
||||
const arr = Array.isArray(inputs[name]) ? inputs[name] : inputs[name].data ?? [...inputs[name]];
|
||||
feeds[name] = new ort.Tensor("int64", BigInt64Array.from(arr.map(Number)), [1, arr.length]);
|
||||
}
|
||||
}
|
||||
|
||||
const { type_logits, bio_logits, age_logits } = await session.run(feeds);
|
||||
|
||||
const typeId = argmax(type_logits.data);
|
||||
const ageId = argmax(age_logits.data);
|
||||
const span = decodeSpanTokens(bio_logits, feeds.attention_mask.data ?? inputs.attention_mask);
|
||||
|
||||
const spanText = span !== null
|
||||
? tokenizer.decode(span.map((i) => Number(inputs.input_ids[i])), { skip_special_tokens: true }).trim()
|
||||
: null;
|
||||
const studio = spanText ? (aliases[spanText.toLowerCase()] ?? spanText) : null;
|
||||
|
||||
console.log(`тип: ${labelMaps.types[typeId]}${ageId === 1 ? " | 18+" : ""}`);
|
||||
console.log(`автор: ${studio ?? "—"}${spanText ? ` (спан «${spanText}»)` : ""}`);
|
||||
console.log(`текст: ${textArg}`);
|
||||
|
||||
// ---------------------------------------------------------------- helpers
|
||||
function argmax(arr) {
|
||||
let best = 0;
|
||||
for (let i = 1; i < arr.length; i++) if (arr[i] > arr[best]) best = i;
|
||||
return best;
|
||||
}
|
||||
|
||||
// первый B(1)-токен и конец I(2)-прогона -> индексы токенов
|
||||
function decodeSpanTokens(bio, attention) {
|
||||
const seq = bio.dims[1];
|
||||
const label = (t) => {
|
||||
const base = t * 3;
|
||||
let bi = 0;
|
||||
if (bio.data[base + 1] > bio.data[base + bi]) bi = 1;
|
||||
if (bio.data[base + 2] > bio.data[base + bi]) bi = 2;
|
||||
return bi;
|
||||
};
|
||||
let s = null;
|
||||
for (let t = 0; t < seq; t++) {
|
||||
if (!attention[t]) break;
|
||||
if (label(t) === 1) { s = t; break; }
|
||||
}
|
||||
if (s === null) return null;
|
||||
let e = s;
|
||||
for (let t = s + 1; t < seq; t++) {
|
||||
if (!attention[t] || label(t) !== 2) break;
|
||||
e = t;
|
||||
}
|
||||
return [s, e];
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
// HTTP-клиент DUB-AI для Node.js (без зависимостей, node >= 18).
|
||||
//
|
||||
// Запуск (требуется запущенный сервис `python -m src.service`):
|
||||
// node predict.mjs --text "Перевод: Дублированный | MovieDalen"
|
||||
// node predict.mjs --html page.html
|
||||
// PREDICT_URL=http://host:8000 node predict.mjs --text "..."
|
||||
const BASE_URL = process.env.PREDICT_URL ?? "http://127.0.0.1:8000";
|
||||
|
||||
function parseArgs(argv) {
|
||||
const args = { text: undefined, html: undefined };
|
||||
for (let i = 2; i < argv.length; i++) {
|
||||
if (argv[i] === "--text") args.text = argv[++i];
|
||||
else if (argv[i] === "--html") args.html = argv[++i];
|
||||
}
|
||||
return args;
|
||||
}
|
||||
|
||||
async function main() {
|
||||
const { text, html } = parseArgs(process.argv);
|
||||
if (Boolean(text) === Boolean(html)) {
|
||||
console.error('нужно ровно одно: --text "..." или --html файл.html');
|
||||
process.exit(1);
|
||||
}
|
||||
const body = text !== undefined ? { text } : { html: await import("node:fs").then(fs => fs.readFileSync(html, "utf-8")) };
|
||||
|
||||
const res = await fetch(`${BASE_URL}/predict`, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(body),
|
||||
});
|
||||
if (!res.ok) {
|
||||
console.error(`HTTP ${res.status}: ${await res.text()}`);
|
||||
process.exit(1);
|
||||
}
|
||||
const { tracks } = await res.json();
|
||||
for (const t of tracks) {
|
||||
console.log(`${t.translation_type}${t.is_18plus ? " 18+" : ""} | ${t.studio ?? "—"} | conf=${t.type_confidence}`);
|
||||
console.log(` текст: ${t.segment}`);
|
||||
if (t.studio_span) console.log(` спан: «${t.studio_span}»`);
|
||||
}
|
||||
}
|
||||
|
||||
main().catch((e) => { console.error(e); process.exit(1); });
|
||||
@@ -0,0 +1,61 @@
|
||||
// Типизированный HTTP-клиент DUB-AI (TypeScript-проекты).
|
||||
// Запуск без компиляции: npx tsx predict.ts --text "Одноголосый закадровый — Дмитрий Есарев"
|
||||
// Зависимости: не требуются (используется глобальный fetch, Node >= 18).
|
||||
|
||||
interface Track {
|
||||
segment: string;
|
||||
translation_type: string;
|
||||
type_confidence: number;
|
||||
studio_span: string | null;
|
||||
studio: string | null;
|
||||
studio_confidence: number | null;
|
||||
is_18plus: boolean;
|
||||
age_confidence: number | null;
|
||||
}
|
||||
|
||||
interface PredictResponse {
|
||||
tracks: Track[];
|
||||
}
|
||||
|
||||
const BASE_URL = process.env.PREDICT_URL ?? "http://127.0.0.1:8000";
|
||||
|
||||
export async function predictOne(text: string): Promise<Track> {
|
||||
const res = await fetch(`${BASE_URL}/predict`, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ text }),
|
||||
});
|
||||
if (!res.ok) throw new Error(`HTTP ${res.status}: ${await res.text()}`);
|
||||
const data: PredictResponse = await res.json();
|
||||
return data.tracks[0];
|
||||
}
|
||||
|
||||
export async function predictHtml(html: string): Promise<Track[]> {
|
||||
const res = await fetch(`${BASE_URL}/predict`, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ html }),
|
||||
});
|
||||
if (!res.ok) throw new Error(`HTTP ${res.status}: ${await res.text()}`);
|
||||
const data: PredictResponse = await res.json();
|
||||
return data.tracks;
|
||||
}
|
||||
|
||||
// --- CLI-демо -------------------------------------------------------------
|
||||
async function main(): Promise<void> {
|
||||
const i = process.argv.indexOf("--text");
|
||||
const j = process.argv.indexOf("--html");
|
||||
if (i !== -1) {
|
||||
const t = await predictOne(process.argv[i + 1]);
|
||||
console.log(`${t.translation_type}${t.is_18plus ? " 18+" : ""} — ${t.studio ?? "—"}`);
|
||||
} else if (j !== -1) {
|
||||
const fs = await import("node:fs");
|
||||
const tracks = await predictHtml(fs.readFileSync(process.argv[j + 1], "utf-8"));
|
||||
for (const t of tracks) console.log(`${t.translation_type}${t.is_18plus ? " 18+" : ""} — ${t.studio ?? "—"}`);
|
||||
} else {
|
||||
console.error('использование: tsx predict.ts --text "..." | --html файл');
|
||||
process.exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
if (process.argv[1]?.endsWith("predict.ts")) main();
|
||||
Reference in new issue
Block a user