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

.//src//data_collection//masakhane_collector.py

@@ -0,0 +1,153 @@
+"""Masakhane benchmark datasets collector."""
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+from datasets import load_dataset
+
+logger = logging.getLogger(__name__)
+
+KINYARWANDA_IDENTIFIERS = ["kin_Latn", "kin", "rw", "kinyarwanda", "kirundi"]
+
+MASAKHANE_DATASETS = [
+    ("masakhane/afrixnli", "train", "kin_Latn"),
+    ("masakhane/afrimmlu", "test", "kin_Latn"),
+    ("masakhane/InjongoIntent", "train", "kin_Latn"),
+    ("masakhane/afrimgsm", "train", "kin_Latn"),
+    ("masakhane/AfriADR", "train", "kin_Latn"),
+    ("masakhane/uhura-arc-easy", "test", "kin_Latn"),
+    ("masakhane/uhura-truthfulqa", "test", "kin_Latn"),
+    ("masakhane/mafand", "train", "kin_Latn"),
+    ("masakhane/afrisenti", "train", "kin_Latn"),
+    ("masakhane/masakhaner2", "train", "kin_Latn"),
+    ("masakhane/masakhaner", "train", "kin_Latn"),
+    ("masakhane/afriqa", "train", "kin_Latn"),
+    ("masakhane/jfleg_kin", "train", None),
+    ("masakhane/masakhapos", "train", "kin_Latn"),
+    ("masakhane/masakhanews", "train", "kin_Latn"),
+    ("masakhane/afri-hate-speech", "train", "kin_Latn"),
+    ("masakhane/uhura-hellaswag", "test", "kin_Latn"),
+    ("masakhane/uhura-winogrande", "test", "kin_Latn"),
+    ("masakhane/uhura-sciq", "test", "kin_Latn"),
+    ("masakhane/uhura-openbookqa", "test", "kin_Latn"),
+    ("masakhane/uhura-mmlu", "test", "kin_Latn"),
+    ("masakhane/uhura-gsm8k", "test", "kin_Latn"),
+    ("masakhane/uhura-arctest", "test", "kin_Latn"),
+    ("facebook/flores", "dev", "kin_Latn"),
+    ("facebook/flores", "devtest", "kin_Latn"),
+    ("openlanguagedata/flores_plus", "dev", "kin_Latn"),
+    ("google/fleurs", "train", "rw"),
+]
+
+
+class MasakhaneCollector:
+    """Collects Masakhane benchmark datasets with Kinyarwanda subsets."""
+
+    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", language: str | None = None) -> dict[str, Any]:
+        """Download and save a Masakhane dataset filtered to Kinyarwanda."""
+        safe_name = name.replace("/", "_")
+        lang_suffix = f"_{language}" if language else ""
+        output_path = self.output_dir / f"{safe_name}{lang_suffix}.jsonl"
+        done_marker = self.output_dir / f"{safe_name}{lang_suffix}.jsonl.done"
+
+        if done_marker.exists():
+            try:
+                row_count = int(done_marker.read_text().strip())
+                logger.info(f"Already collected: {name} (lang={language}, {row_count} rows), skipping")
+                return {
+                    "name": name,
+                    "split": split,
+                    "language": language,
+                    "rows": row_count,
+                    "output_path": str(output_path),
+                    "status": "success",
+                }
+            except (ValueError, OSError):
+                pass
+
+        logger.info(f"Collecting Masakhane dataset: {name} (lang={language})")
+
+        try:
+            kwargs = {"split": split}
+            if language:
+                kwargs["name"] = language
+
+            dataset = load_dataset(name, **kwargs)
+
+            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
+
+            done_marker.write_text(str(count))
+
+            return {
+                "name": name,
+                "split": split,
+                "language": language,
+                "rows": count,
+                "output_path": str(output_path),
+                "status": "success",
+            }
+
+        except Exception as e:
+            logger.error(f"Failed to collect {name} (lang={language}): {e}")
+
+            if language and language in KINYARWANDA_IDENTIFIERS:
+                logger.info(f"  Retrying {name} without language filter...")
+                try:
+                    dataset = load_dataset(name, split=split)
+                    count = 0
+                    with open(output_path, "w", encoding="utf-8") as f:
+                        for item in dataset:
+                            item_str = json.dumps(item, ensure_ascii=False)
+                            if any(kw in item_str.lower() for kw in ["kinyarwanda", "kin_", "rwanda"]):
+                                f.write(item_str + "\n")
+                                count += 1
+
+                    done_marker.write_text(str(count))
+                    return {
+                        "name": name,
+                        "split": split,
+                        "language": language,
+                        "rows": count,
+                        "output_path": str(output_path),
+                        "status": "success",
+                    }
+                except Exception as e2:
+                    logger.error(f"  Retry also failed for {name}: {e2}")
+
+            return {"name": name, "status": "failed", "error": str(e)}
+
+    def collect_all(self) -> list[dict[str, Any]]:
+        """Collect all Masakhane benchmarks with Kinyarwanda filtering."""
+        results = []
+        for name, split, lang in MASAKHANE_DATASETS:
+            result = self.collect_dataset(name, split, lang)
+            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": len(results),
+            "successful": successful,
+            "failed": failed,
+            "total_rows": total_rows,
+        }
+
+        summary_path = self.output_dir / "masakhane_summary.json"
+        with open(summary_path, "w", encoding="utf-8") as f:
+            json.dump(summary, f, indent=2, ensure_ascii=False)
+
+        logger.info(f"Masakhane collection: {summary}")
+        return results