mii/tokenizers
1
0
Fork 0
mirror of https://github.com/mii443/tokenizers.git synced 2025-08-23 00:35:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

Also accept iterators of batches in train_from_iterator

Browse Source
This commit is contained in:
Anthony MOI
2020-11-24 22:48:55 -05:00
committed by Anthony MOI
parent e0a70f1fb2
commit 75deaecdd0
3 changed files with 25 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
bindings/python/src/tokenizer.rs
View File

@ -1104,15 +1104,30 @@ impl PyTokenizer {
.map(|_| {})
});
ResultShunt::process(iterator.map(|seq| seq?.extract::<&str>()), |iter| {
if let Some(send) = sender.take() {
for seq in iter {
send.send(seq)
.map_err(|e| exceptions::PyException::new_err(e.to_string()))?;
ResultShunt::process(
// Each element of the iterator can either be:
// - An iterator, to allow batching
// - A string
iterator.flat_map(|seq| match seq {
Ok(s) => {
if let Ok(iter) = s.iter() {
itertools::Either::Left(iter.map(|i| i?.extract::<&str>()))
} else {
itertools::Either::Right(std::iter::once(s.extract::<&str>()))
}
}
}
Ok(())
})?
Err(e) => itertools::Either::Right(std::iter::once(Err(e))),
}),
|iter| {
if let Some(send) = sender.take() {
for seq in iter {
send.send(seq)
.map_err(|e| exceptions::PyException::new_err(e.to_string()))?;
}
}
Ok(())
},
)?
})
.unwrap()
}