Upload .\src\data_collection\parallel_corpus_collector.py with huggingface_hub

.//src//data_collection//parallel_corpus_collector.py

@@ -0,0 +1,111 @@
+"""Parallel corpus collector for massive multilingual datasets."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+from datasets import load_dataset
+
+logger = logging.getLogger(__name__)
+
+
+class ParallelCorpusCollector:
+    """Collects parallel translation corpora from HuggingFace."""
+
+    def __init__(self, output_dir: str, config: dict[str, Any]):
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self.config = config
+
+    def collect_dataset(self, name: str, split: str = "train", max_rows: int | None = None) -> dict[str, Any]:
+        """Download and save a parallel dataset using streaming."""
+        safe_name = name.replace("/", "_")
+        output_path = self.output_dir / f"{safe_name}.jsonl"
+        done_marker = self.output_dir / f"{safe_name}.jsonl.done"
+
+        if done_marker.exists():
+            try:
+                row_count = int(done_marker.read_text().strip())
+                logger.info(f"Already collected: {name} ({row_count} rows), skipping")
+                return {
+                    "name": name,
+                    "split": split,
+                    "rows": row_count,
+                    "output_path": str(output_path),
+                    "status": "success",
+                }
+            except (ValueError, OSError):
+                pass
+
+        logger.info(f"Collecting parallel corpus: {name} (split: {split})")
+
+        try:
+            dataset = load_dataset(name, split=split, streaming=True)
+
+            count = 0
+            with open(output_path, "w", encoding="utf-8") as f:
+                for item in dataset:
+                    f.write(json.dumps(item, ensure_ascii=False) + "\n")
+                    count += 1
+                    if max_rows and count >= max_rows:
+                        break
+                    if count % 100_000 == 0:
+                        logger.info(f"  {name}: {count:,} rows collected...")
+
+            done_marker.write_text(str(count))
+
+            return {
+                "name": name,
+                "split": split,
+                "rows": count,
+                "output_path": str(output_path),
+                "status": "success",
+            }
+
+        except Exception as e:
+            logger.error(f"Failed to collect {name}: {e}")
+            return {"name": name, "split": split, "status": "failed", "error": str(e)}
+
+    def collect_all(self) -> list[dict[str, Any]]:
+        """Collect all parallel corpora from config."""
+        datasets_config = self.config.get("huggingface", {}).get("datasets", [])
+
+        parallel_keywords = ["sentence-pairs", "parallel", "bitext", "ccmatrix", "nllb"]
+        parallel_datasets = []
+        seen_names = set()
+
+        for ds in datasets_config:
+            name = ds.get("name", "")
+            desc = ds.get("description", "").lower()
+            if name in seen_names:
+                continue
+            if any(kw in name.lower() or kw in desc for kw in parallel_keywords):
+                if "michsethowusu" in name or "sentence-pairs" in name.lower():
+                    parallel_datasets.append(name)
+                    seen_names.add(name)
+
+        logger.info(f"Found {len(parallel_datasets)} parallel datasets from config")
+
+        results = []
+        for name in parallel_datasets:
+            result = self.collect_dataset(name)
+            results.append(result)
+
+        successful = sum(1 for r in results if r["status"] == "success")
+        failed = sum(1 for r in results if r["status"] == "failed")
+        total_rows = sum(r.get("rows", 0) for r in results if r["status"] == "success")
+
+        summary = {
+            "total_datasets": len(results),
+            "successful": successful,
+            "failed": failed,
+            "total_rows": total_rows,
+        }
+
+        summary_path = self.output_dir / "parallel_summary.json"
+        with open(summary_path, "w", encoding="utf-8") as f:
+            json.dump(summary, f, indent=2, ensure_ascii=False)
+
+        logger.info(f"Parallel corpus collection: {summary}")
+        return results