New version. Staticmethods need to return a IntoPy<PyObject>

which is non trivial for PyClassInitializer. Instead I added a lower
staticmethod that returns raw objects, and the `from_file(s)` methods
are implemented directly in Python.
This commit is contained in:
Nicolas Patry
2020-09-22 18:10:56 +02:00
parent 98a30eead1
commit 60c1e25910
8 changed files with 72 additions and 31 deletions

View File

@@ -80,6 +80,10 @@ class BertWordPieceTokenizer(BaseTokenizer):
super().__init__(tokenizer, parameters)
def from_file(vocab: str, **kwargs):
vocab = WordPiece.read_file(vocab)
return BertWordPieceTokenizer(vocab, **kwargs)
def train(
self,
files: Union[str, List[str]],

View File

@@ -77,6 +77,11 @@ class ByteLevelBPETokenizer(BaseTokenizer):
super().__init__(tokenizer, parameters)
@staticmethod
def from_files(vocab_filename: str, merges_filename: str, **kwargs):
vocab, merges = BPE.read_files(vocab_filename, merges_filename)
return ByteLevelBPETokenizer(vocab, merges, **kwargs)
def train(
self,
files: Union[str, List[str]],

View File

@@ -94,6 +94,11 @@ class CharBPETokenizer(BaseTokenizer):
super().__init__(tokenizer, parameters)
@staticmethod
def from_files(vocab_filename: str, merges_filename: str, **kwargs):
vocab, merges = BPE.read_files(vocab_filename, merges_filename)
return CharBPETokenizer(vocab, merges, **kwargs)
def train(
self,
files: Union[str, List[str]],

View File

@@ -47,6 +47,11 @@ class SentencePieceBPETokenizer(BaseTokenizer):
super().__init__(tokenizer, parameters)
@staticmethod
def from_files(vocab_filename: str, merges_filename: str, **kwargs):
vocab, merges = BPE.read_files(vocab_filename, merges_filename)
return SentencePieceBPETokenizer(vocab, merges, **kwargs)
def train(
self,
files: Union[str, List[str]],

View File

@@ -62,8 +62,6 @@ class BPE(Model):
fuse_unk: (`optional`) bool:
Multiple unk tokens get fused into only 1
"""
@staticmethod
def __init__(
self,
vocab: Optional[Union[str, Dict[str, int]]],
@@ -77,6 +75,15 @@ class BPE(Model):
):
pass
@staticmethod
def read_files(vocab_filename: str, merges_filename: str) -> Tuple[Vocab, Merges]:
pass
@staticmethod
def from_files(vocab_filename: str, merges_filename: str, **kwargs) -> BPE:
vocab, merges = BPE.read_files(vocab_filename, merges_filename)
return BPE(vocab, merges, **kwargs)
class WordPiece(Model):
""" WordPiece model class
@@ -101,6 +108,15 @@ class WordPiece(Model):
):
pass
@staticmethod
def read_file(vocab_filename: str) -> Tuple[Vocab]:
pass
@staticmethod
def from_files(vocab_filename: str, **kwargs) -> WordPiece:
vocab = WordPiece.read_files(vocab_filename)
return WordPiece(vocab, **kwargs)
class WordLevel(Model):
"""
Most simple tokenizer model based on mapping token from a vocab file to their corresponding id.
@@ -118,6 +134,15 @@ class WordLevel(Model):
def __init__(self, vocab: Optional[Union[str, Dict[str, int]]], unk_token: Optional[str]):
pass
@staticmethod
def read_file(vocab_filename: str) -> Tuple[Vocab]:
pass
@staticmethod
def from_files(vocab_filename: str, **kwargs) -> WordLevel:
vocab = WordLevel.read_files(vocab_filename)
return WordLevel(vocab, **kwargs)
class Unigram(Model):
"""UnigramEncoding model class

View File

@@ -7,7 +7,7 @@ use pyo3::exceptions;
use pyo3::prelude::*;
use pyo3::types::*;
use serde::{Deserialize, Serialize};
use tk::models::bpe::{BpeBuilder, BPE};
use tk::models::bpe::{BpeBuilder, Merges, Vocab, BPE};
use tk::models::unigram::Unigram;
use tk::models::wordlevel::WordLevel;
use tk::models::wordpiece::{WordPiece, WordPieceBuilder};
@@ -217,16 +217,13 @@ impl PyBPE {
}
#[staticmethod]
#[args(kwargs = "**")]
fn from_files(
vocab_filename: String,
merges_filename: String,
kwargs: Option<&PyDict>,
) -> PyResult<(Self, PyModel)> {
let mut builder = BPE::builder();
builder = builder.files(vocab_filename, merges_filename);
PyBPE::with_builder(builder, kwargs)
fn read_files(vocab_filename: &str, merges_filename: &str) -> PyResult<(Vocab, Merges)> {
BPE::read_files(vocab_filename, merges_filename).map_err(|e| {
exceptions::PyValueError::new_err(format!(
"Error while reading vocab&merges files: {}",
e
))
})
}
}
@@ -292,10 +289,10 @@ impl PyWordPiece {
}
#[staticmethod]
fn from_file(vocab: String, kwargs: Option<&PyDict>) -> PyResult<(Self, PyModel)> {
let mut builder = WordPiece::builder();
builder = builder.files(vocab);
PyWordPiece::with_builder(builder, kwargs)
fn read_file(vocab_filename: &str) -> PyResult<Vocab> {
WordPiece::read_files(vocab_filename).map_err(|e| {
exceptions::PyValueError::new_err(format!("Error while reading WordPiece file: {}", e))
})
}
}
@@ -356,15 +353,10 @@ impl PyWordLevel {
}
#[staticmethod]
fn from_file(vocab_filename: &str, kwargs: Option<&PyDict>) -> PyResult<(Self, PyModel)> {
let unk_token = PyWordLevel::get_unk(kwargs)?;
let model = WordLevel::from_files(vocab_filename, unk_token).map_err(|e| {
exceptions::PyException::new_err(format!(
"Error while loading WordLevel from file: {}",
e
))
})?;
Ok((PyWordLevel {}, PyModel::new(Arc::new(model.into()))))
fn read_file(vocab_filename: &str) -> PyResult<Vocab> {
WordLevel::read_files(vocab_filename).map_err(|e| {
exceptions::PyValueError::new_err(format!("Error while reading WordLevel file: {}", e))
})
}
}