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

.//src//data_collection//huggingface_collector.py

@@ -0,0 +1,172 @@
+"""Huggingface dataset collector for Kinyarwanda data."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+from datasets import load_dataset
+
+logger = logging.getLogger(__name__)
+
+MAX_ROWS_DEFAULT = 2_000_000_000
+
+DATASET_MAX_ROWS = {
+    "HuggingFaceFW/fineweb-2": 500_000_000,
+    "HPLT/HPLT2.0_cleaned": 100_000_000,
+    "yhavinga/ccmatrix": 100_000_000,
+    "sentence-transformers/parallel-sentences-ccmatrix": 50_000_000,
+    "datalama/pretrain-nllb-filtered": 50_000_000,
+    "hotchpotch/nllb-english-bitext-hq": 50_000_000,
+    "NaolBM/african-corpus": 10_000_000,
+    "bonadossou/afrolm_active_learning_dataset": 5_000_000,
+    "wikimedia/wikidata-title-desc": 10_000_000,
+    "wikimedia/wikipedia": 5_000_000,
+    "csebuetnlp/xlsum": 2_000_000,
+    "csebuetnlp/mt5_xlsum_all_languages_44": 1_000_000,
+}
+
+
+class HuggingfaceCollector:
+    """Collects Kinyarwanda datasets 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
+        self.datasets_info: list[dict[str, Any]] = []
+
+    def collect_dataset(self, dataset_config: dict[str, str]) -> dict[str, Any]:
+        """Download and save a single dataset using streaming."""
+        name = dataset_config["name"]
+        split = dataset_config.get("split", "train")
+
+        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),
+                    "description": dataset_config.get("description", ""),
+                    "status": "success",
+                }
+            except (ValueError, OSError):
+                pass
+
+        logger.info(f"Collecting dataset: {name} (split: {split})")
+
+        max_rows = DATASET_MAX_ROWS.get(name, MAX_ROWS_DEFAULT)
+
+        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 count >= max_rows:
+                        logger.info(f"  Reached max_rows limit ({max_rows}) for {name}")
+                        break
+                    if count % 100_000 == 0:
+                        logger.info(f"  {name}: {count:,} rows collected...")
+
+            done_marker.write_text(str(count))
+
+            info = {
+                "name": name,
+                "split": split,
+                "rows": count,
+                "output_path": str(output_path),
+                "description": dataset_config.get("description", ""),
+                "status": "success",
+            }
+
+            logger.info(f"  Saved {count:,} rows to {output_path}")
+            self.datasets_info.append(info)
+            return info
+
+        except Exception as e:
+            logger.error(f"  Failed to collect {name}: {e}")
+            info = {
+                "name": name,
+                "split": split,
+                "status": "failed",
+                "error": str(e),
+                "description": dataset_config.get("description", ""),
+            }
+            self.datasets_info.append(info)
+            return info
+
+    def collect_all(self) -> list[dict[str, Any]]:
+        """Collect all configured datasets."""
+        datasets_config = self.config.get("huggingface", {}).get("datasets", [])
+
+        logger.info(f"Collecting {len(datasets_config)} datasets from Huggingface...")
+
+        for dataset_config in datasets_config:
+            self.collect_dataset(dataset_config)
+
+        summary_path = self.output_dir / "huggingface_summary.json"
+        with open(summary_path, "w", encoding="utf-8") as f:
+            json.dump(self.datasets_info, f, indent=2, ensure_ascii=False)
+
+        successful = sum(1 for d in self.datasets_info if d["status"] == "success")
+        failed = sum(1 for d in self.datasets_info if d["status"] == "failed")
+        total_rows = sum(d.get("rows", 0) for d in self.datasets_info if d["status"] == "success")
+
+        logger.info(f"Collection complete: {successful} succeeded, {failed} failed, {total_rows:,} total rows")
+
+        return self.datasets_info
+
+    def get_text_samples(self, max_samples: int = 1000) -> list[str]:
+        """Extract text samples from all collected datasets."""
+        texts = []
+
+        for info in self.datasets_info:
+            if info["status"] != "success":
+                continue
+
+            output_path = Path(info["output_path"])
+            if not output_path.exists():
+                continue
+
+            with open(output_path, "r", encoding="utf-8") as f:
+                for i, line in enumerate(f):
+                    if i >= max_samples:
+                        break
+
+                    item = json.loads(line)
+
+                    text = self._extract_text(item)
+                    if text:
+                        texts.append(text)
+
+        return texts
+
+    def _extract_text(self, item: dict[str, Any]) -> str | None:
+        """Extract text content from a dataset item."""
+        text_keys = ["text", "content", "sentence", "document", "paragraph", "translation"]
+
+        for key in text_keys:
+            if key in item:
+                value = item[key]
+                if isinstance(value, str):
+                    return value
+                elif isinstance(value, dict):
+                    for subkey in ["rw", "rwanda", "kinyarwanda", "kin"]:
+                        if subkey in value:
+                            return value[subkey]
+
+        for key in item:
+            if isinstance(item[key], str) and len(item[key]) > 50:
+                return item[key]
+
+        return None