mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
110 lines
3.2 KiB
Python
110 lines
3.2 KiB
Python
from collections import Counter
|
|
|
|
from datasets import Dataset
|
|
|
|
from ctx_to_lora.data.preprocessing_fn import get_preprocessing_fn
|
|
from ctx_to_lora.data.processing import (
|
|
OOLONG_MAX_CONTEXT_LEN,
|
|
filter_dataset_by_max_value,
|
|
stratified_subsample_dataset,
|
|
)
|
|
from ctx_to_lora.eval_utils import compute_oolong_metrics, extract_oolong_answer
|
|
|
|
|
|
def test_oolong_preprocessing_maps_core_fields_and_metadata():
|
|
preprocess = get_preprocessing_fn("oolong-synth", is_eval=True)
|
|
sample = {
|
|
"id": 7,
|
|
"dataset": "spam",
|
|
"context_window_text": "context body",
|
|
"question": "Give your final answer in the form 'Label: answer'.",
|
|
"task_group": "classification",
|
|
"task": "most_freq",
|
|
"answer": "['ham']",
|
|
"answer_type": "ANSWER_TYPE.LABEL",
|
|
"input_subset": "all",
|
|
"num_labels": 2,
|
|
"context_window_id": 101,
|
|
}
|
|
|
|
processed = preprocess(sample)
|
|
|
|
assert processed["context"] == "context body"
|
|
assert processed["prompts"] == [sample["question"]]
|
|
assert processed["responses"] == ["['ham']"]
|
|
assert processed["oolong_task"] == "most_freq"
|
|
assert processed["oolong_answer_type"] == "ANSWER_TYPE.LABEL"
|
|
assert processed["oolong_context_window_id"] == 101
|
|
|
|
|
|
def test_extract_oolong_answer_prefers_suffix_after_colon():
|
|
answer, confidence = extract_oolong_answer("Reasoning...\nLabel: [ham]")
|
|
|
|
assert answer == "ham"
|
|
assert confidence == "vhigh"
|
|
|
|
|
|
def test_compute_oolong_metrics_handles_exact_numeric_and_date_scoring():
|
|
decoded_txts = [
|
|
{
|
|
"generated": "Label: ham",
|
|
"label": "['ham']",
|
|
"oolong_answer_type": "ANSWER_TYPE.LABEL",
|
|
},
|
|
{
|
|
"generated": "Answer: 8",
|
|
"label": "['10']",
|
|
"oolong_answer_type": "ANSWER_TYPE.NUMERIC",
|
|
},
|
|
{
|
|
"generated": "Date: 2025-02-07",
|
|
"label": "[datetime.date(2025, 2, 7)]",
|
|
"oolong_answer_type": "ANSWER_TYPE.DATE",
|
|
},
|
|
]
|
|
|
|
metrics, per_sample, counts = compute_oolong_metrics(decoded_txts)
|
|
|
|
assert metrics["oolong_score"] == (1.0 + (0.75**2) + 1.0) / 3
|
|
assert metrics["oolong_parse_rate"] == 1.0
|
|
assert per_sample["oolong_score"] == [1.0, 0.75**2, 1.0]
|
|
assert counts["oolong_score"] == 3
|
|
|
|
|
|
def test_stratified_subsample_dataset_balances_context_len_groups():
|
|
ds = Dataset.from_dict(
|
|
{
|
|
"context_len": [1] * 800 + [2] * 50 + [3] * 800,
|
|
"row_id": list(range(1650)),
|
|
}
|
|
)
|
|
|
|
sampled = stratified_subsample_dataset(
|
|
ds=ds,
|
|
group_column="context_len",
|
|
max_samples=1000,
|
|
seed=42,
|
|
)
|
|
counts = Counter(sampled["context_len"])
|
|
|
|
assert len(sampled) == 1000
|
|
assert counts[2] == 50
|
|
assert counts[1] == 475
|
|
assert counts[3] == 475
|
|
|
|
|
|
def test_filter_dataset_by_max_value_drops_over_64k_context_len_rows():
|
|
ds = Dataset.from_dict(
|
|
{
|
|
"context_len": [1024, OOLONG_MAX_CONTEXT_LEN, OOLONG_MAX_CONTEXT_LEN + 1],
|
|
"row_id": [0, 1, 2],
|
|
}
|
|
)
|
|
|
|
filtered = filter_dataset_by_max_value(
|
|
ds=ds,
|
|
column="context_len",
|
|
max_value=OOLONG_MAX_CONTEXT_LEN,
|
|
)
|
|
|
|
assert filtered["context_len"] == [1024, OOLONG_MAX_CONTEXT_LEN]
|