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

.//src//data_collection//wikipedia_collector.py

@@ -0,0 +1,261 @@
+"""Wikipedia data collector for Kinyarwanda."""
+
+import bz2
+import json
+import logging
+import xml.etree.ElementTree as ET
+from pathlib import Path
+from typing import Any
+
+import requests
+from tqdm import tqdm
+
+logger = logging.getLogger(__name__)
+
+WIKI_API = "https://rw.wikipedia.org/w/api.php"
+
+
+class WikipediaCollector:
+    """Collects Kinyarwanda Wikipedia content."""
+
+    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.session = requests.Session()
+        self.session.headers.update({"User-Agent": "BwengeAi/0.1 (Kinyarwanda AI Research)"})
+
+    def get_site_stats(self) -> dict[str, Any]:
+        """Get Wikipedia site statistics."""
+        params = {
+            "action": "query",
+            "meta": "siteinfo",
+            "siprop": "statistics",
+            "format": "json",
+        }
+        
+        response = self.session.get(WIKI_API, params=params)
+        response.raise_for_status()
+        data = response.json()
+        
+        stats = data.get("query", {}).get("statistics", {})
+        logger.info(f"Kinyarwanda Wikipedia stats: {stats}")
+        return stats
+
+    def fetch_article_list(self) -> list[str]:
+        """Fetch list of all article titles."""
+        articles = []
+        params = {
+            "action": "query",
+            "list": "allpages",
+            "apnamespace": 0,
+            "aplimit": "max",
+            "format": "json",
+        }
+        
+        while True:
+            response = self.session.get(WIKI_API, params=params)
+            response.raise_for_status()
+            data = response.json()
+            
+            pages = data.get("query", {}).get("allpages", [])
+            articles.extend(page["title"] for page in pages)
+            
+            if "continue" in data:
+                params["apcontinue"] = data["continue"]["apcontinue"]
+            else:
+                break
+        
+        logger.info(f"Found {len(articles)} articles")
+        return articles
+
+    def fetch_article_content(self, title: str) -> dict[str, Any] | None:
+        """Fetch content for a single article."""
+        params = {
+            "action": "parse",
+            "page": title,
+            "prop": "wikitext|text|categories|links",
+            "format": "json",
+        }
+        
+        try:
+            response = self.session.get(WIKI_API, params=params)
+            response.raise_for_status()
+            data = response.json()
+            
+            parse = data.get("parse", {})
+            text = parse.get("text", {}).get("*", "")
+            clean_text = self._clean_html(text)
+            
+            return {
+                "title": title,
+                "text": clean_text,
+                "wikitext": parse.get("wikitext", {}).get("*", ""),
+                "categories": [cat["*"] for cat in parse.get("categories", [])],
+                "links": [link["*"] for link in parse.get("links", [])],
+                "length": len(clean_text),
+            }
+        except Exception as e:
+            logger.warning(f"Failed to fetch article '{title}': {e}")
+            return None
+
+    def _clean_html(self, html: str) -> str:
+        """Clean HTML content to plain text."""
+        from bs4 import BeautifulSoup
+        
+        soup = BeautifulSoup(html, "lxml")
+        
+        for tag in soup(["script", "style", "sup", "sub"]):
+            tag.decompose()
+        
+        text = soup.get_text(separator=" ", strip=True)
+        
+        lines = text.split()
+        clean_lines = []
+        for line in lines:
+            line = line.strip()
+            if line and not line.startswith("[") and len(line) > 2:
+                clean_lines.append(line)
+        
+        return " ".join(clean_lines)
+
+    def collect_via_api(self, max_articles: int | None = None) -> list[dict[str, Any]]:
+        """Collect articles via the Wikipedia API."""
+        articles = self.fetch_article_list()
+        
+        if max_articles:
+            articles = articles[:max_articles]
+        
+        logger.info(f"Collecting {len(articles)} articles via API...")
+        
+        results = []
+        output_path = self.output_dir / "wikipedia_articles.jsonl"
+        
+        with open(output_path, "w", encoding="utf-8") as f:
+            for title in tqdm(articles, desc="Fetching Wikipedia articles"):
+                article = self.fetch_article_content(title)
+                if article and article["text"]:
+                    f.write(json.dumps(article, ensure_ascii=False) + "\n")
+                    results.append(article)
+        
+        summary = {
+            "source": "kinyarwanda_wikipedia",
+            "total_articles": len(articles),
+            "collected": len(results),
+            "output_path": str(output_path),
+            "total_characters": sum(r["length"] for r in results),
+        }
+        
+        summary_path = self.output_dir / "wikipedia_summary.json"
+        with open(summary_path, "w", encoding="utf-8") as f:
+            json.dump(summary, f, indent=2, ensure_ascii=False)
+        
+        logger.info(f"Collected {len(results)} articles ({summary['total_characters']} characters)")
+        return results
+
+    def download_dump(self) -> Path | None:
+        """Download the Wikipedia dump file."""
+        dump_config = self.config.get("wikipedia", {})
+        dump_url = dump_config.get("dump_url", "https://dumps.wikimedia.org/rwwiki/latest/")
+        dump_file = dump_config.get("dump_file", "rwwiki-latest-pages-articles.xml.bz2")
+        
+        url = f"{dump_url}{dump_file}"
+        output_path = self.output_dir / dump_file
+        
+        if output_path.exists():
+            logger.info(f"Dump file already exists: {output_path}")
+            return output_path
+        
+        logger.info(f"Downloading dump from {url}...")
+        
+        try:
+            response = requests.get(url, stream=True, timeout=300)
+            response.raise_for_status()
+            
+            total_size = int(response.headers.get("content-length", 0))
+            
+            with open(output_path, "wb") as f:
+                with tqdm(total=total_size, unit="B", unit_scale=True, desc="Downloading") as pbar:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+            
+            logger.info(f"Downloaded dump to {output_path}")
+            return output_path
+            
+        except Exception as e:
+            logger.error(f"Failed to download dump: {e}")
+            return None
+
+    def parse_dump(self, dump_path: Path, max_articles: int | None = None) -> list[dict[str, Any]]:
+        """Parse a Wikipedia XML dump file using iterative XML parsing."""
+        logger.info(f"Parsing dump file: {dump_path}")
+
+        results = []
+        count = 0
+        title = ""
+        in_page = False
+        in_text = False
+        text_parts: list[str] = []
+
+        with bz2.open(dump_path, "rb") as f:
+            for event, elem in ET.iterparse(f, events=("start", "end")):
+                tag = elem.tag.split("}")[-1] if "}" in elem.tag else elem.tag
+
+                if event == "start" and tag == "page":
+                    in_page = True
+                    title = ""
+                    text_parts = []
+
+                elif event == "start" and tag == "title" and in_page:
+                    title = (elem.text or "").strip()
+
+                elif event == "start" and tag == "text" and in_page:
+                    in_text = True
+                    if elem.text:
+                        text_parts.append(elem.text)
+
+                elif event == "end" and tag == "text" and in_page:
+                    in_text = False
+
+                elif event == "end" and tag == "page" and in_page:
+                    in_page = False
+                    full_text = " ".join(text_parts)
+                    clean_text = self._clean_mediawiki(full_text)
+
+                    if clean_text and len(clean_text) > 100:
+                        results.append({
+                            "title": title,
+                            "text": clean_text,
+                            "length": len(clean_text),
+                        })
+                        count += 1
+
+                        if count % 1000 == 0:
+                            logger.info(f"  Parsed {count:,} articles...")
+
+                        if max_articles and count >= max_articles:
+                            break
+
+                    elem.clear()
+
+        output_path = self.output_dir / "wikipedia_dump_articles.jsonl"
+        with open(output_path, "w", encoding="utf-8") as f:
+            for article in results:
+                f.write(json.dumps(article, ensure_ascii=False) + "\n")
+
+        logger.info(f"Parsed {len(results)} articles from dump")
+        return results
+
+    def _clean_mediawiki(self, text: str) -> str:
+        """Clean MediaWiki markup to plain text."""
+        import re
+        
+        text = re.sub(r"\[\[([^|\]]*\|)?([^\]]*)\]\]", r"\2", text)
+        text = re.sub(r"\{\{[^}]*\}\}", "", text)
+        text = re.sub(r"'''?", "", text)
+        text = re.sub(r"<[^>]+>", "", text)
+        text = re.sub(r"==+\s*[^=]*\s*==+", "", text)
+        text = re.sub(r"\s+", " ", text)
+        
+        return text.strip()