Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add deserialize for pre tokenizers #1603

Merged
merged 7 commits into from
Aug 8, 2024
Merged

add deserialize for pre tokenizers #1603

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
da394ee
add deserialize
ArthurZucker Aug 7, 2024
04e8032
copy from the decoder
ArthurZucker Aug 7, 2024
794209a
fmt
ArthurZucker Aug 7, 2024
e99daff
clippy
ArthurZucker Aug 7, 2024
dd6652b
fix rust tests
ArthurZucker Aug 7, 2024
f275e03
fmt
ArthurZucker Aug 7, 2024
0fe7b80
don't change the test
ArthurZucker Aug 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
153 changes: 143 additions & 10 deletions tokenizers/src/pre_tokenizers/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ pub mod split;
pub mod unicode_scripts;
pub mod whitespace;

use serde::{Deserialize, Serialize};
use serde::{Deserialize, Deserializer, Serialize};

use crate::pre_tokenizers::bert::BertPreTokenizer;
use crate::pre_tokenizers::byte_level::ByteLevel;
Expand All @@ -23,7 +23,7 @@ use crate::pre_tokenizers::unicode_scripts::UnicodeScripts;
use crate::pre_tokenizers::whitespace::{Whitespace, WhitespaceSplit};
use crate::{PreTokenizedString, PreTokenizer};

#[derive(Deserialize, Serialize, Clone, Debug, PartialEq)]
#[derive(Serialize, Clone, Debug, PartialEq)]
#[serde(untagged)]
pub enum PreTokenizerWrapper {
BertPreTokenizer(BertPreTokenizer),
Expand Down Expand Up @@ -57,6 +57,142 @@ impl PreTokenizer for PreTokenizerWrapper {
}
}

impl<'de> Deserialize<'de> for PreTokenizerWrapper {
fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Deserialize)]
pub struct Tagged {
#[serde(rename = "type")]
variant: EnumType,
#[serde(flatten)]
rest: serde_json::Value,
}
#[derive(Deserialize, Serialize)]
pub enum EnumType {
BertPreTokenizer,
ByteLevel,
Delimiter,
Metaspace,
Whitespace,
Sequence,
Split,
Punctuation,
WhitespaceSplit,
Digits,
UnicodeScripts,
}

#[derive(Deserialize)]
#[serde(untagged)]
pub enum PreTokenizerHelper {
Tagged(Tagged),
Legacy(serde_json::Value),
}

#[derive(Deserialize)]
#[serde(untagged)]
pub enum PreTokenizerUntagged {
BertPreTokenizer(BertPreTokenizer),
ByteLevel(ByteLevel),
Delimiter(CharDelimiterSplit),
Metaspace(Metaspace),
Whitespace(Whitespace),
Sequence(Sequence),
Split(Split),
Punctuation(Punctuation),
WhitespaceSplit(WhitespaceSplit),
Digits(Digits),
UnicodeScripts(UnicodeScripts),
}

let helper = PreTokenizerHelper::deserialize(deserializer)?;

Ok(match helper {
PreTokenizerHelper::Tagged(pretok) => {
let mut values: serde_json::Map<String, serde_json::Value> =
serde_json::from_value(pretok.rest).map_err(serde::de::Error::custom)?;
values.insert(
"type".to_string(),
serde_json::to_value(&pretok.variant).map_err(serde::de::Error::custom)?,
);
let values = serde_json::Value::Object(values);
match pretok.variant {
EnumType::BertPreTokenizer => PreTokenizerWrapper::BertPreTokenizer(
serde_json::from_value(values).map_err(serde::de::Error::custom)?,
),
EnumType::ByteLevel => PreTokenizerWrapper::ByteLevel(
serde_json::from_value(values).map_err(serde::de::Error::custom)?,
),
EnumType::Delimiter => PreTokenizerWrapper::Delimiter(
serde_json::from_value(values).map_err(serde::de::Error::custom)?,
),
EnumType::Metaspace => PreTokenizerWrapper::Metaspace(
serde_json::from_value(values).map_err(serde::de::Error::custom)?,
),
EnumType::Whitespace => PreTokenizerWrapper::Whitespace(
serde_json::from_value(values).map_err(serde::de::Error::custom)?,
),
EnumType::Sequence => PreTokenizerWrapper::Sequence(
serde_json::from_value(values).map_err(serde::de::Error::custom)?,
),
EnumType::Split => PreTokenizerWrapper::Split(
serde_json::from_value(values).map_err(serde::de::Error::custom)?,
),
EnumType::Punctuation => PreTokenizerWrapper::Punctuation(
serde_json::from_value(values).map_err(serde::de::Error::custom)?,
),
EnumType::WhitespaceSplit => PreTokenizerWrapper::WhitespaceSplit(
serde_json::from_value(values).map_err(serde::de::Error::custom)?,
),
EnumType::Digits => PreTokenizerWrapper::Digits(
serde_json::from_value(values).map_err(serde::de::Error::custom)?,
),
EnumType::UnicodeScripts => PreTokenizerWrapper::UnicodeScripts(
serde_json::from_value(values).map_err(serde::de::Error::custom)?,
),
}
}

PreTokenizerHelper::Legacy(value) => {
let untagged = serde_json::from_value(value).map_err(serde::de::Error::custom)?;
match untagged {
PreTokenizerUntagged::BertPreTokenizer(bert) => {
PreTokenizerWrapper::BertPreTokenizer(bert)
}
PreTokenizerUntagged::ByteLevel(byte_level) => {
PreTokenizerWrapper::ByteLevel(byte_level)
}
PreTokenizerUntagged::Delimiter(delimiter) => {
PreTokenizerWrapper::Delimiter(delimiter)
}
PreTokenizerUntagged::Metaspace(metaspace) => {
PreTokenizerWrapper::Metaspace(metaspace)
}
PreTokenizerUntagged::Whitespace(whitespace) => {
PreTokenizerWrapper::Whitespace(whitespace)
}
PreTokenizerUntagged::Sequence(sequence) => {
PreTokenizerWrapper::Sequence(sequence)
}
PreTokenizerUntagged::Split(split) => PreTokenizerWrapper::Split(split),
PreTokenizerUntagged::Punctuation(punctuation) => {
PreTokenizerWrapper::Punctuation(punctuation)
}
PreTokenizerUntagged::WhitespaceSplit(whitespace_split) => {
PreTokenizerWrapper::WhitespaceSplit(whitespace_split)
}
PreTokenizerUntagged::Digits(digits) => PreTokenizerWrapper::Digits(digits),
PreTokenizerUntagged::UnicodeScripts(unicode_scripts) => {
PreTokenizerWrapper::UnicodeScripts(unicode_scripts)
}
}
}
})
}
}

impl_enum_from!(BertPreTokenizer, PreTokenizerWrapper, BertPreTokenizer);
impl_enum_from!(ByteLevel, PreTokenizerWrapper, ByteLevel);
impl_enum_from!(CharDelimiterSplit, PreTokenizerWrapper, Delimiter);
Expand Down Expand Up @@ -152,33 +288,30 @@ mod tests {
match reconstructed {
Err(err) => assert_eq!(
err.to_string(),
"data did not match any variant of untagged enum PreTokenizerWrapper"
"data did not match any variant of untagged enum PreTokenizerUntagged"
),
_ => panic!("Expected an error here"),
}

let json = r#"{"type":"Metaspace", "replacement":"▁" }"#;
let reconstructed = serde_json::from_str::<PreTokenizerWrapper>(json);
let reconstructed = serde_json::from_str::<PreTokenizerWrapper>(json).unwrap();
assert_eq!(
reconstructed.unwrap(),
reconstructed,
PreTokenizerWrapper::Metaspace(Metaspace::default())
);

let json = r#"{"type":"Metaspace", "add_prefix_space":true }"#;
let reconstructed = serde_json::from_str::<PreTokenizerWrapper>(json);
match reconstructed {
Err(err) => assert_eq!(
err.to_string(),
"data did not match any variant of untagged enum PreTokenizerWrapper"
),
Err(err) => assert_eq!(err.to_string(), "missing field `replacement`"),
_ => panic!("Expected an error here"),
}
let json = r#"{"behavior":"default_split"}"#;
let reconstructed = serde_json::from_str::<PreTokenizerWrapper>(json);
match reconstructed {
Err(err) => assert_eq!(
err.to_string(),
"data did not match any variant of untagged enum PreTokenizerWrapper"
"data did not match any variant of untagged enum PreTokenizerUntagged"
),
_ => panic!("Expected an error here"),
}
Expand Down
Loading