wehub-resource-sync
119 lines
3.6 KiB
Python
119 lines
3.6 KiB
Python
import pandas as pd
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data.preprocessors import Tokenizer
|
|
|
|
|
|
def test_tokenizer():
|
|
"""Tests basic Tokenizer functionality."""
|
|
|
|
col_a = ["this is a test", "apple"]
|
|
col_b = ["the quick brown fox jumps over the lazy dog", "banana banana"]
|
|
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
|
ds = ray.data.from_pandas(in_df)
|
|
|
|
tokenizer = Tokenizer(["A", "B"])
|
|
transformed = tokenizer.transform(ds)
|
|
out_df = transformed.to_pandas()
|
|
|
|
processed_col_a = [["this", "is", "a", "test"], ["apple"]]
|
|
processed_col_b = [
|
|
["the", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog"],
|
|
["banana", "banana"],
|
|
]
|
|
expected_df = pd.DataFrame.from_dict(
|
|
{"A": processed_col_a, "B": processed_col_b}
|
|
).astype(out_df.dtypes.to_dict())
|
|
|
|
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
|
|
|
# Test append mode
|
|
with pytest.raises(
|
|
ValueError, match="The length of columns and output_columns must match."
|
|
):
|
|
Tokenizer(columns=["A", "B"], output_columns=["A_tokenized"])
|
|
|
|
tokenizer = Tokenizer(
|
|
columns=["A", "B"], output_columns=["A_tokenized", "B_tokenized"]
|
|
)
|
|
transformed = tokenizer.transform(ds)
|
|
out_df = transformed.to_pandas()
|
|
print(out_df)
|
|
expected_df = pd.DataFrame.from_dict(
|
|
{
|
|
"A": col_a,
|
|
"B": col_b,
|
|
"A_tokenized": processed_col_a,
|
|
"B_tokenized": processed_col_b,
|
|
}
|
|
).astype(out_df.dtypes.to_dict())
|
|
|
|
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
|
|
|
# Test custom tokenization function
|
|
def custom_tokenizer(s: str) -> list:
|
|
return s.replace("banana", "fruit").split()
|
|
|
|
tokenizer = Tokenizer(
|
|
columns=["A", "B"],
|
|
tokenization_fn=custom_tokenizer,
|
|
output_columns=["A_custom", "B_custom"],
|
|
)
|
|
transformed = tokenizer.transform(ds)
|
|
out_df = transformed.to_pandas()
|
|
|
|
custom_processed_col_a = [["this", "is", "a", "test"], ["apple"]]
|
|
custom_processed_col_b = [
|
|
["the", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog"],
|
|
["fruit", "fruit"],
|
|
]
|
|
expected_df = pd.DataFrame.from_dict(
|
|
{
|
|
"A": col_a,
|
|
"B": col_b,
|
|
"A_custom": custom_processed_col_a,
|
|
"B_custom": custom_processed_col_b,
|
|
}
|
|
).astype(out_df.dtypes.to_dict())
|
|
|
|
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
|
|
|
|
|
def test_tokenizer_serialization():
|
|
"""Test Tokenizer serialization and deserialization functionality."""
|
|
from ray.data.preprocessor import SerializablePreprocessorBase
|
|
|
|
# Create tokenizer
|
|
tokenizer = Tokenizer(columns=["text"])
|
|
|
|
# Serialize using CloudPickle
|
|
serialized = tokenizer.serialize()
|
|
|
|
# Verify it's binary CloudPickle format
|
|
assert isinstance(serialized, bytes)
|
|
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
|
|
|
# Deserialize
|
|
deserialized = Tokenizer.deserialize(serialized)
|
|
|
|
# Verify type and field values
|
|
assert isinstance(deserialized, Tokenizer)
|
|
assert deserialized.columns == ["text"]
|
|
assert callable(deserialized.tokenization_fn)
|
|
assert deserialized.output_columns == ["text"]
|
|
|
|
# Verify it works correctly
|
|
df = pd.DataFrame({"text": ["hello world", "foo bar"]})
|
|
result = deserialized.transform_batch(df)
|
|
|
|
# Verify tokenization was applied correctly
|
|
assert result["text"][0] == ["hello", "world"]
|
|
assert result["text"][1] == ["foo", "bar"]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-sv", __file__]))
|