Fix deadlocks with custom python components.

This commit is contained in:
Sebastian Pütz
2020-08-03 13:34:38 +02:00
committed by Anthony MOI
parent 88e6eb5db6
commit 27e326ab2b
2 changed files with 26 additions and 16 deletions

View File

@ -583,12 +583,15 @@ impl Tokenizer {
Ok(input)
})
.collect::<PyResult<Vec<tk::EncodeInput>>>()?;
ToPyResult(
self.tokenizer
.encode_batch(input, add_special_tokens)
.map(|encodings| encodings.into_iter().map(|e| e.into()).collect()),
)
.into()
let gil = Python::acquire_gil();
gil.python().allow_threads(|| {
ToPyResult(
self.tokenizer
.encode_batch(input, add_special_tokens)
.map(|encodings| encodings.into_iter().map(|e| e.into()).collect()),
)
.into()
})
}
fn decode(&self, ids: Vec<u32>, skip_special_tokens: Option<bool>) -> PyResult<String> {
@ -604,11 +607,14 @@ impl Tokenizer {
sentences: Vec<Vec<u32>>,
skip_special_tokens: Option<bool>,
) -> PyResult<Vec<String>> {
ToPyResult(
self.tokenizer
.decode_batch(sentences, skip_special_tokens.unwrap_or(true)),
)
.into()
let gil = Python::acquire_gil();
gil.python().allow_threads(|| {
ToPyResult(
self.tokenizer
.decode_batch(sentences, skip_special_tokens.unwrap_or(true)),
)
.into()
})
}
fn token_to_id(&self, token: &str) -> Option<u32> {
@ -661,11 +667,14 @@ impl Tokenizer {
fn train(&mut self, trainer: &Trainer, files: Vec<String>) -> PyResult<()> {
trainer.trainer.execute(|trainer| {
if let Err(e) = self.tokenizer.train(trainer, files) {
Err(exceptions::Exception::py_err(format!("{}", e)))
} else {
Ok(())
}
let gil = Python::acquire_gil();
gil.python().allow_threads(|| {
if let Err(e) = self.tokenizer.train(trainer, files) {
Err(exceptions::Exception::py_err(format!("{}", e)))
} else {
Ok(())
}
})
})
}