Spaces:
Sleeping
Sleeping
File size: 5,919 Bytes
9521b3b e84e505 c8685ee 9521b3b c8685ee 9521b3b c8685ee 9521b3b c8685ee 9521b3b c8685ee 9521b3b c8685ee 9521b3b c8685ee 9521b3b c8685ee 9521b3b c8685ee 9521b3b c8685ee 9521b3b c8685ee 9521b3b c8685ee 9521b3b c8685ee 9521b3b c8685ee 9521b3b c8685ee e84e505 9521b3b c8685ee e84e505 | 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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 | """Phase 2 data cleaning and splitting.
Loads raw questions, removes duplicates, and writes train/val/test splits.
"""
import json
import os
import sys
import pandas as pd
from sklearn.model_selection import train_test_split
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import config # noqa: E402
class DatasetProcessor:
"""Clean, deduplicate, and split the raw dataset."""
def __init__(self, processing_config):
"""Store processing config."""
self.cfg = processing_config
def load_raw(self):
"""Load raw JSONL into a DataFrame."""
records = []
with open(self.cfg.raw_path, "r", encoding="utf-8") as handle:
for line in handle:
line = line.strip()
if line:
records.append(json.loads(line))
frame = pd.DataFrame(records)
print(f"Loaded {len(frame)} raw questions.")
return frame
@staticmethod
def _normalise(text):
"""Normalize text for exact duplicate checks."""
return " ".join(text.lower().split())
def remove_exact_duplicates(self, frame):
"""Drop exact duplicates and keep the first one."""
before = len(frame)
frame = frame.copy()
frame["_norm"] = frame["question"].map(self._normalise)
frame = frame.drop_duplicates(subset="_norm", keep="first")
frame = frame.drop(columns="_norm").reset_index(drop=True)
print(f"Exact dedup: {before} -> {len(frame)} "
f"({before - len(frame)} removed).")
return frame
def remove_near_duplicates(self, frame):
"""Remove near-duplicate questions with embedding similarity."""
from sentence_transformers import SentenceTransformer, util
before = len(frame)
questions = frame["question"].tolist()
print(f"Encoding {before} questions for near-dup detection "
f"(model: {self.cfg.embed_model})...")
model = SentenceTransformer(self.cfg.embed_model)
embeddings = model.encode(
questions, convert_to_tensor=True, show_progress_bar=True
)
# Returns [score, i, j] sorted by score.
pairs = util.paraphrase_mining_embeddings(embeddings)
removed = set()
for score, i, j in pairs:
if score < self.cfg.near_dup_threshold:
break
low, high = (i, j) if i < j else (j, i)
if low not in removed and high not in removed:
removed.add(high)
kept_mask = [idx not in removed for idx in range(before)]
frame = frame[kept_mask].reset_index(drop=True)
print(f"Near-dup removal (threshold {self.cfg.near_dup_threshold}): "
f"{before} -> {len(frame)} ({before - len(frame)} removed).")
return frame
def split_in_domain(self, frame):
"""Split in-domain rows into train/val/test."""
in_domain = frame[frame["split"] == "train"].copy().reset_index(drop=True)
ood = frame[frame["split"] == "ood"].copy().reset_index(drop=True)
strat_key = in_domain["domain"] + "|" + in_domain["bloom_class"]
# Split off test first.
train_val, test = train_test_split(
in_domain,
test_size=self.cfg.test_size,
stratify=strat_key,
random_state=self.cfg.random_state,
)
# Then split validation from the rest.
strat_key_tv = train_val["domain"] + "|" + train_val["bloom_class"]
relative_val = self.cfg.val_size / (1.0 - self.cfg.test_size)
train, val = train_test_split(
train_val,
test_size=relative_val,
stratify=strat_key_tv,
random_state=self.cfg.random_state,
)
splits = {
"train": train.reset_index(drop=True),
"val": val.reset_index(drop=True),
"test": test.reset_index(drop=True),
"ood_test": ood,
}
return splits
def save_splits(self, splits):
"""Save each split as a JSONL file."""
os.makedirs(self.cfg.processed_dir, exist_ok=True)
keep_cols = [
"question", "bloom_sublevel", "bloom_class",
"domain", "topic", "source",
]
for name, frame in splits.items():
path = os.path.join(self.cfg.processed_dir, f"{name}.jsonl")
with open(path, "w", encoding="utf-8") as handle:
for _, row in frame.iterrows():
record = {col: row[col] for col in keep_cols if col in row}
handle.write(json.dumps(record, ensure_ascii=False) + "\n")
print(f" wrote {len(frame):4d} -> {path}")
@staticmethod
def print_breakdown(splits):
"""Print split sizes and class/domain counts."""
print("\nSplit sizes:")
for name, frame in splits.items():
print(f" {name:9s} {len(frame)}")
print("\nClass balance per split:")
for name, frame in splits.items():
counts = frame["bloom_class"].value_counts().to_dict()
print(f" {name:9s} {counts}")
print("\nPer-(domain x class) counts (all splits combined):")
combined = pd.concat(splits.values())
table = combined.groupby(["domain", "bloom_class"]).size().unstack(fill_value=0)
print(table.to_string())
def run(self):
"""Run the full Phase 2 pipeline."""
frame = self.load_raw()
frame = self.remove_exact_duplicates(frame)
frame = self.remove_near_duplicates(frame)
splits = self.split_in_domain(frame)
self.save_splits(splits)
self.print_breakdown(splits)
print("\nPhase 2 complete.")
def main():
"""Run Phase 2 processing."""
processor = DatasetProcessor(config.ProcessingConfig())
processor.run()
if __name__ == "__main__":
main()
|