Spaces:
Sleeping
Sleeping
| 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 | |
| } | |