Spaces:
Sleeping
Sleeping
File size: 1,652 Bytes
e71990c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | import numpy as np
import pandas as pd
import pytest
import ray
from src.data import clean_text, stratify_split, tokenize
@pytest.mark.parametrize(
"text, expected",
[
("Hello, World!", "hello world"),
("Stocks rally after the Fed meeting", "stocks rally fed meeting"),
("Check https://example.com for details", "check details"),
("", ""),
],
)
def test_clean_text(text, expected):
assert clean_text(text) == expected
def test_clean_text_is_lowercase():
assert clean_text("UPPER CASE TEXT") == clean_text("UPPER CASE TEXT").lower()
def test_stratify_split_preserves_class_proportions():
df = pd.DataFrame(
{
"text": [f"sample {i}" for i in range(100)],
"label": [0] * 50 + [1] * 50,
}
)
ds = ray.data.from_pandas(df)
train_ds, test_ds = stratify_split(ds, stratify="label", test_size=0.2)
train_df = train_ds.to_pandas()
test_df = test_ds.to_pandas()
assert len(train_df) + len(test_df) == len(df)
assert set(train_df["label"].unique()) == {0, 1}
assert set(test_df["label"].unique()) == {0, 1}
assert train_df["label"].value_counts()[0] == train_df["label"].value_counts()[1]
assert test_df["label"].value_counts()[0] == test_df["label"].value_counts()[1]
def test_tokenize_shapes_match_batch_size():
batch = pd.DataFrame({"text": ["hello world", "a longer sentence here"], "label": [0, 1]})
result = tokenize(batch)
assert result["ids"].shape[0] == len(batch)
assert result["masks"].shape == result["ids"].shape
assert np.array_equal(result["targets"], np.array([0, 1]))
|