diff --git a/src/args.rs b/src/args.rs index 5de5cf5..da0aa2c 100644 --- a/src/args.rs +++ b/src/args.rs @@ -394,6 +394,10 @@ pub struct Arg { #[arg(short = 'L', long, global = true, default_value_t = 5, value_parser = parse_flac_compression_level)] /// FLAC compression level for output FLAC audio files. 0 means fastest compression, 8 means best compression. pub flac_compression_level: u32, + #[arg(long, global = true)] + /// Add a mark to the end of each message for LLM translation. + /// Only works on m3t format. + pub llm_trans_mark: Option, #[command(subcommand)] /// Command pub command: Command, diff --git a/src/main.rs b/src/main.rs index 030f2ff..71f6055 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1101,7 +1101,10 @@ pub fn import_script( continue; } }; - let mut parser = output_scripts::m3t::M3tParser::new(&s); + let mut parser = output_scripts::m3t::M3tParser::new( + &s, + arg.llm_trans_mark.as_ref().map(|s| s.as_str()), + ); match parser.parse() { Ok(mes) => mes, Err(e) => { @@ -1320,7 +1323,10 @@ pub fn import_script( let enc = get_output_encoding(arg); let b = utils::files::read_file(&out_f)?; let s = utils::encoding::decode_to_string(enc, &b, true)?; - let mut parser = output_scripts::m3t::M3tParser::new(&s); + let mut parser = output_scripts::m3t::M3tParser::new( + &s, + arg.llm_trans_mark.as_ref().map(|s| s.as_str()), + ); parser.parse()? } types::OutputScriptType::Yaml => { diff --git a/src/output_scripts/m3t.rs b/src/output_scripts/m3t.rs index abadca4..a3f0016 100644 --- a/src/output_scripts/m3t.rs +++ b/src/output_scripts/m3t.rs @@ -15,12 +15,17 @@ use anyhow::Result; pub struct M3tParser<'a> { str: &'a str, line: usize, + llm_mark: Option<&'a str>, } impl<'a> M3tParser<'a> { /// Creates a new M3tParser with the given string. - pub fn new(str: &'a str) -> Self { - M3tParser { str, line: 1 } + pub fn new(str: &'a str, llm_mark: Option<&'a str>) -> Self { + M3tParser { + str, + line: 1, + llm_mark, + } } fn next_line(&mut self) -> Option<&'a str> { @@ -68,16 +73,31 @@ impl<'a> M3tParser<'a> { .is_empty() { llm.take() + .map(|s| { + let mut s = s.to_string(); + if let Some(mark) = self.llm_mark { + s.push_str(mark); + } + s + }) .unwrap_or_else(|| { - if message.starts_with("「") { + String::from(if message.starts_with("「") { "「」" } else { "" - } + }) }) .replace("\\n", "\n") } else { - message.replace("\\n", "\n") + let mut tmp = message.replace("\\n", "\n"); + if let Some(llm) = llm.take() { + if tmp == llm { + if let Some(mark) = self.llm_mark { + tmp.push_str(mark); + } + } + } + tmp }; messages.push(Message::new(message, name.take())); } else {