DUB-AI: multi-task модель (тип озвучки + автор-спан + 18+), обучение/инференс/ONNX/сервис, тесты и документация

examples/rust/onnx_local/src/main.rs

@@ -0,0 +1,140 @@
+//! Локальный инференс DUB-AI без Python (путь Б: ONNX).
+//!
+//! cargo run --release -- <путь к models/best-onnx> "Перевод: Дублированный | MovieDalen"
+//!
+//! Повторяет препроцессинг/постобработку src/predict.py:
+//!   токенизация -> ONNX (type_logits, bio_logits, age_logits) ->
+//!   argmax типа, 18+ и первого B-I-прогона по offsets.
+
+use std::collections::HashMap;
+
+use ort::session::builder::GraphOptimizationLevel;
+use ort::session::Session;
+use ort::value::Tensor;
+use tokenizers::Tokenizer;
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let args: Vec<String> = std::env::args().collect();
+    if args.len() < 3 {
+        eprintln!("использование: cargo run --release -- <models/best-onnx> \"текст дорожки\"");
+        std::process::exit(1);
+    }
+    let model_dir = &args[1];
+    let text = &args[2];
+
+    // --- метки и алиасы ---------------------------------------------------
+    let maps: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(
+        format!("{model_dir}/label_maps.json"),
+    )?)?;
+    let types: Vec<String> = maps["types"]
+        .as_array()
+        .unwrap()
+        .iter()
+        .map(|v| v.as_str().unwrap().to_string())
+        .collect();
+    let aliases: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(
+        format!("{model_dir}/aliases.json"),
+    )?)?;
+    let alias_map: HashMap<String, String> = aliases["aliases"]
+        .as_object()
+        .unwrap()
+        .iter()
+        .map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string()))
+        .collect();
+
+    // --- токенизация -------------------------------------------------------
+    let tok = Tokenizer::from_file(format!("{model_dir}/tokenizer.json"))
+        .map_err(|e| e.to_string())?;
+    let enc = tok.encode(text.as_str(), true).map_err(|e| e.to_string())?;
+    let ids: Vec<i64> = enc.get_ids().iter().map(|&v| v as i64).collect();
+    let mask: Vec<i64> = enc.get_attention_mask().iter().map(|&v| v as i64).collect();
+    let type_ids: Vec<i64> = enc.get_type_ids().iter().map(|&v| v as i64).collect();
+    let offsets = enc.get_offsets().to_vec();
+    let seq = ids.len();
+
+    // --- ONNX-инференс ------------------------------------------------------
+    let env = ort::init().build()?;
+    let mut session = Session::builder(&env)?
+        .with_optimization_level(GraphOptimizationLevel::Level3)?
+        .with_intra_threads(2)?
+        .commit_from_file(format!("{model_dir}/model.onnx"))?;
+
+    let shape = [1usize, seq];
+    let outputs = session.run(ort::inputs![
+        "input_ids" => Tensor::from_array((shape, ids))?,
+        "attention_mask" => Tensor::from_array((shape, mask.clone()))?,
+        "token_type_ids" => Tensor::from_array((shape, type_ids))?,
+    ])?;
+
+    let (_, type_logits) = outputs["type_logits"].try_extract_tensor::<f32>()?;
+    let (_, age_logits) = outputs["age_logits"].try_extract_tensor::<f32>()?;
+    let (_, bio_flat) = outputs["bio_logits"].try_extract_tensor::<f32>()?;
+
+    let type_id = argmax(type_logits.iter().copied());
+    let age_id = argmax(age_logits.iter().copied());
+    let label = |token: usize| -> u32 {
+        let base = token * 3;
+        let (mut bi, mut bv) = (0u32, bio_flat[base]);
+        for (c, &v) in bio_flat[base + 1..base + 3].iter().enumerate() {
+            if v > bv {
+                bi = c as u32 + 1;
+                bv = v;
+            }
+        }
+        bi
+    };
+
+    // спан: первый B-токен + цепочка I (только реальные токены)
+    let mut start_tok = None;
+    for t in 0..seq {
+        if mask[t] == 0 {
+            break;
+        }
+        if offsets[t].0 >= offsets[t].1 {
+            continue; // спецтокены
+        }
+        if label(t) == 1 {
+            start_tok = Some(t);
+            break;
+        }
+    }
+    let mut span_txt: Option<String> = None;
+    if let Some(s) = start_tok {
+        let mut e = s;
+        for t in s + 1..seq {
+            if mask[t] == 0 || offsets[t].0 >= offsets[t].1 || label(t) != 2 {
+                break;
+            }
+            e = t;
+        }
+        span_txt = Some(text[offsets[s].0..offsets[e].1].trim().to_string());
+    }
+    let studio = span_txt.as_ref().map(|s| {
+        alias_map
+            .get(&s.to_lowercase())
+            .cloned()
+            .unwrap_or_else(|| s.clone())
+    });
+
+    println!(
+        "тип: {}{}",
+        types[type_id],
+        if age_id == 1 { "  |  18+" } else { "" }
+    );
+    match (&studio, &span_txt) {
+        (Some(s), Some(sp)) => println!("автор: {s}  (спан «{sp}»)"),
+        _ => println!("автор: —"),
+    }
+    println!("текст: {text}");
+    Ok(())
+}
+
+fn argmax<I: Iterator<Item = f32>>(it: I) -> usize {
+    let mut best = (0usize, f32::MIN);
+    for (i, v) in it.enumerate() {
+        if v > best.1 {
+            best = (i, v);
+        }
+    }
+    best.0
+}