Spaces:
Sleeping
Sleeping
File size: 3,099 Bytes
c1b1f3a | 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 | import os
import json
import hashlib
from typing import List, Dict, Tuple, Any
import pandas as pd
from sklearn.model_selection import train_test_split
from .schema import DocumentRecord
def compute_file_hash(filepath: str) -> str:
hasher = hashlib.sha256()
with open(filepath, "rb") as f:
while chunk := f.read(8192):
hasher.update(chunk)
return hasher.hexdigest()
def build_stratified_splits(
records: List[DocumentRecord],
splits_dir: str,
processed_dir: str,
seed: int = 42,
train_ratio: float = 0.70,
val_ratio: float = 0.15,
test_ratio: float = 0.15
) -> Dict[str, Any]:
"""
Creates stratified 70/15/15 train/val/test splits, writes manifests as CSV,
saves unified records to jsonl, and returns metadata including SHA256 hashes
and class weight balancing strategy.
"""
os.makedirs(splits_dir, exist_ok=True)
os.makedirs(processed_dir, exist_ok=True)
jsonl_path = os.path.join(processed_dir, "unified_documents.jsonl")
with open(jsonl_path, "w", encoding="utf-8") as f_out:
for r in records:
f_out.write(r.model_dump_json() + "\n")
# Build DataFrame for stratification
df_meta = pd.DataFrame([
{
"doc_id": r.doc_id,
"source_dataset": r.source_dataset,
"doc_class": r.doc_class,
"image_path": r.image_path,
"ocr_token_count": len(r.ocr_words)
}
for r in records
])
# Stratified split: train vs temp (val + test)
temp_ratio = val_ratio + test_ratio
train_df, temp_df = train_test_split(
df_meta,
test_size=temp_ratio,
random_state=seed,
stratify=df_meta["doc_class"]
)
# Split temp into val and test
val_relative_ratio = val_ratio / temp_ratio
val_df, test_df = train_test_split(
temp_df,
test_size=(1.0 - val_relative_ratio),
random_state=seed,
stratify=temp_df["doc_class"]
)
train_path = os.path.join(splits_dir, "train.csv")
val_path = os.path.join(splits_dir, "val.csv")
test_path = os.path.join(splits_dir, "test.csv")
train_df.to_csv(train_path, index=False)
val_df.to_csv(val_path, index=False)
test_df.to_csv(test_path, index=False)
# Class distribution and balanced weights
class_counts = train_df["doc_class"].value_counts().to_dict()
total_samples = len(train_df)
n_classes = len(class_counts)
class_weights = {
cls: total_samples / (n_classes * count)
for cls, count in class_counts.items()
}
manifest_hashes = {
"train_manifest_hash": compute_file_hash(train_path),
"val_manifest_hash": compute_file_hash(val_path),
"test_manifest_hash": compute_file_hash(test_path),
"jsonl_hash": compute_file_hash(jsonl_path)
}
return {
"train_count": len(train_df),
"val_count": len(val_df),
"test_count": len(test_df),
"class_counts_train": class_counts,
"class_weights": class_weights,
"hashes": manifest_hashes
}
|