regen.mde 0.8.0 → 0.8.2

package/tools/collect_online_docx_corpus.py

@@ -1,272 +1,272 @@
-from __future__ import annotations
-
-import argparse
-import hashlib
-import json
-import re
-import sys
-import time
-from pathlib import Path
-from typing import Iterable
-from urllib.parse import parse_qs, quote_plus, unquote, urlparse
-from zipfile import BadZipFile, ZipFile
-
-import requests
-
-
-DEFAULT_QUERIES = [
-    'filetype:docx "equation"',
-    'filetype:docx "equations"',
-    'filetype:docx "Cambria Math"',
-    'filetype:docx "math equations"',
-    'filetype:docx "Microsoft Equation"',
-    'filetype:docx "OMML"',
-    'filetype:docx "quadratic equation"',
-    'filetype:docx "integral"',
-    'filetype:docx "matrix"',
-    'filetype:docx "calculus"',
-    'filetype:docx "physics" "equation"',
-    'filetype:docx "engineering" "equation"',
-    'filetype:docx "statistics" "equation"',
-    'filetype:docx "algebra" "equation"',
-    'filetype:docx "geometry" "equation"',
-]
-
-GITHUB_REPO_QUERIES = [
-    "docx equation math",
-    "docx equations latex",
-    "docx omml latex",
-    "word equation docx",
-    "markdown docx equation",
-    "pandoc docx equation",
-    "math docx converter",
-    "equation-heavy docx",
-    "ooxml math docx",
-    "Cambria Math docx",
-]
-
-HEADERS = {
-    "User-Agent": (
-        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
-        "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0 Safari/537.36"
-    )
-}
-
-DOCX_CT_HINTS = (
-    "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
-    "application/octet-stream",
-    "application/zip",
-    "binary/octet-stream",
-)
-
-
-def extract_urls(html: str) -> list[str]:
-    raw = re.findall(r'https?://[^"\'<>\s)]+', html)
-    urls: list[str] = []
-    for url in raw:
-        url = unquote(url).replace("&amp;", "&")
-        parsed = urlparse(url)
-        if parsed.netloc.endswith("bing.com") and parsed.path == "/ck/a":
-            qs = parse_qs(parsed.query)
-            for key in ("u", "r"):
-                if key in qs:
-                    url = unquote(qs[key][0])
-                    break
-        if "duckduckgo.com/l/?" in url:
-            qs = parse_qs(urlparse(url).query)
-            if "uddg" in qs:
-                url = unquote(qs["uddg"][0])
-        urls.append(url.rstrip(".,;"))
-    return urls
-
-
-def search(query: str, pages: int = 2) -> Iterable[str]:
-    endpoints = [
-        "https://www.bing.com/search?q={query}&first={offset}",
-        "https://html.duckduckgo.com/html/?q={query}&s={offset}",
-    ]
-    for endpoint in endpoints:
-        for page in range(pages):
-            offset = page * 10 + 1
-            url = endpoint.format(query=quote_plus(query), offset=offset)
-            try:
-                response = requests.get(url, headers=HEADERS, timeout=20)
-                if response.status_code >= 400:
-                    continue
-                yield from extract_urls(response.text)
-                time.sleep(0.4)
-            except requests.RequestException:
-                continue
-
-
-def github_json(url: str) -> dict | None:
-    try:
-        response = requests.get(url, headers={**HEADERS, "Accept": "application/vnd.github+json"}, timeout=25)
-        if response.status_code >= 400:
-            return None
-        return response.json()
-    except requests.RequestException:
-        return None
-
-
-def github_docx_urls(max_repos_per_query: int = 8) -> Iterable[str]:
-    seen_repos: set[str] = set()
-    for query in GITHUB_REPO_QUERIES:
-        search_url = (
-            "https://api.github.com/search/repositories"
-            f"?q={quote_plus(query)}&per_page={max_repos_per_query}"
-        )
-        payload = github_json(search_url)
-        if not payload:
-            continue
-        for repo in payload.get("items", []):
-            full_name = repo.get("full_name")
-            branch = repo.get("default_branch") or "main"
-            if not full_name or full_name in seen_repos:
-                continue
-            seen_repos.add(full_name)
-            tree_url = f"https://api.github.com/repos/{full_name}/git/trees/{branch}?recursive=1"
-            tree = github_json(tree_url)
-            if not tree:
-                continue
-            for item in tree.get("tree", []):
-                path = item.get("path", "")
-                if item.get("type") == "blob" and path.lower().endswith(".docx"):
-                    yield f"https://raw.githubusercontent.com/{full_name}/{branch}/{quote_path(path)}"
-            time.sleep(0.25)
-
-
-def quote_path(path: str) -> str:
-    return "/".join(quote_plus(part).replace("+", "%20") for part in path.split("/"))
-
-
-def looks_like_docx_url(url: str) -> bool:
-    low = url.lower()
-    return ".docx" in low and not any(bad in low for bad in ("?format=pdf", "/view?", "webcache"))
-
-
-def safe_name(index: int, url: str, content: bytes) -> str:
-    parsed = urlparse(url)
-    stem = Path(unquote(parsed.path)).name or f"online-{index:03d}.docx"
-    stem = re.sub(r"[^A-Za-z0-9._ -]+", "_", stem)
-    if not stem.lower().endswith(".docx"):
-        stem += ".docx"
-    digest = hashlib.sha256(content).hexdigest()[:10]
-    return f"{index:03d}-{digest}-{stem}"
-
-
-def validate_docx(path: Path) -> dict:
-    with ZipFile(path) as zf:
-        names = set(zf.namelist())
-        if "[Content_Types].xml" not in names or "word/document.xml" not in names:
-            raise BadZipFile("not a Word DOCX package")
-        document_xml = zf.read("word/document.xml")
-        return {
-            "omml_nodes": len(re.findall(rb"<m:oMath\b|<m:oMathPara\b", document_xml)),
-            "media_parts": len([name for name in names if name.startswith("word/media/")]),
-            "document_xml_bytes": len(document_xml),
-        }
-
-
-def download(url: str, out_dir: Path, index: int, max_mb: int) -> dict | None:
-    try:
-        with requests.get(url, headers=HEADERS, timeout=30, stream=True, allow_redirects=True) as response:
-            if response.status_code >= 400:
-                return None
-            content_type = response.headers.get("content-type", "").split(";")[0].strip().lower()
-            if content_type and content_type not in DOCX_CT_HINTS and "wordprocessingml" not in content_type:
-                if ".docx" not in response.url.lower():
-                    return None
-            chunks: list[bytes] = []
-            total = 0
-            limit = max_mb * 1024 * 1024
-            for chunk in response.iter_content(chunk_size=1024 * 256):
-                if not chunk:
-                    continue
-                total += len(chunk)
-                if total > limit:
-                    return None
-                chunks.append(chunk)
-        content = b"".join(chunks)
-    except requests.RequestException:
-        return None
-
-    if not content.startswith(b"PK"):
-        return None
-
-    out_path = out_dir / safe_name(index, response.url, content)
-    out_path.write_bytes(content)
-    try:
-        stats = validate_docx(out_path)
-    except (BadZipFile, KeyError):
-        out_path.unlink(missing_ok=True)
-        return None
-
-    return {
-        "file": str(out_path),
-        "source_url": response.url,
-        "bytes": len(content),
-        "content_type": content_type,
-        **stats,
-    }
-
-
-def main() -> int:
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--out", type=Path, required=True)
-    parser.add_argument("--target", type=int, default=50)
-    parser.add_argument("--max-mb", type=int, default=40)
-    parser.add_argument("--prefer-equations", action="store_true")
-    args = parser.parse_args()
-
-    args.out.mkdir(parents=True, exist_ok=True)
-    seen: set[str] = set()
-    manifest: list[dict] = []
-    index = len(list(args.out.glob("*.docx"))) + 1
-
-    print("source: github repositories", flush=True)
-    for url in github_docx_urls():
-        if len(manifest) >= args.target:
-            break
-        normalized = url.split("#", 1)[0]
-        if normalized in seen:
-            continue
-        seen.add(normalized)
-        item = download(normalized, args.out, index, args.max_mb)
-        if not item:
-            continue
-        manifest.append(item)
-        index += 1
-        print(f"kept {len(manifest):02d}: {Path(item['file']).name} omml={item['omml_nodes']}", flush=True)
-
-    for query in DEFAULT_QUERIES:
-        print(f"search: {query}", flush=True)
-        for url in search(query):
-            if len(manifest) >= args.target:
-                break
-            if not looks_like_docx_url(url):
-                continue
-            normalized = url.split("#", 1)[0]
-            if normalized in seen:
-                continue
-            seen.add(normalized)
-            item = download(normalized, args.out, index, args.max_mb)
-            if not item:
-                continue
-            if args.prefer_equations and item["omml_nodes"] == 0 and len(manifest) < args.target // 2:
-                Path(item["file"]).unlink(missing_ok=True)
-                continue
-            manifest.append(item)
-            index += 1
-            print(f"kept {len(manifest):02d}: {Path(item['file']).name} omml={item['omml_nodes']}", flush=True)
-        if len(manifest) >= args.target:
-            break
-
-    manifest_path = args.out / "online-docx-manifest.json"
-    manifest_path.write_text(json.dumps({"count": len(manifest), "items": manifest}, indent=2), encoding="utf-8")
-    print(f"saved manifest: {manifest_path}", flush=True)
-    return 0 if len(manifest) >= args.target else 1
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import re
+import sys
+import time
+from pathlib import Path
+from typing import Iterable
+from urllib.parse import parse_qs, quote_plus, unquote, urlparse
+from zipfile import BadZipFile, ZipFile
+
+import requests
+
+
+DEFAULT_QUERIES = [
+    'filetype:docx "equation"',
+    'filetype:docx "equations"',
+    'filetype:docx "Cambria Math"',
+    'filetype:docx "math equations"',
+    'filetype:docx "Microsoft Equation"',
+    'filetype:docx "OMML"',
+    'filetype:docx "quadratic equation"',
+    'filetype:docx "integral"',
+    'filetype:docx "matrix"',
+    'filetype:docx "calculus"',
+    'filetype:docx "physics" "equation"',
+    'filetype:docx "engineering" "equation"',
+    'filetype:docx "statistics" "equation"',
+    'filetype:docx "algebra" "equation"',
+    'filetype:docx "geometry" "equation"',
+]
+
+GITHUB_REPO_QUERIES = [
+    "docx equation math",
+    "docx equations latex",
+    "docx omml latex",
+    "word equation docx",
+    "markdown docx equation",
+    "pandoc docx equation",
+    "math docx converter",
+    "equation-heavy docx",
+    "ooxml math docx",
+    "Cambria Math docx",
+]
+
+HEADERS = {
+    "User-Agent": (
+        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+        "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0 Safari/537.36"
+    )
+}
+
+DOCX_CT_HINTS = (
+    "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+    "application/octet-stream",
+    "application/zip",
+    "binary/octet-stream",
+)
+
+
+def extract_urls(html: str) -> list[str]:
+    raw = re.findall(r'https?://[^"\'<>\s)]+', html)
+    urls: list[str] = []
+    for url in raw:
+        url = unquote(url).replace("&amp;", "&")
+        parsed = urlparse(url)
+        if parsed.netloc.endswith("bing.com") and parsed.path == "/ck/a":
+            qs = parse_qs(parsed.query)
+            for key in ("u", "r"):
+                if key in qs:
+                    url = unquote(qs[key][0])
+                    break
+        if "duckduckgo.com/l/?" in url:
+            qs = parse_qs(urlparse(url).query)
+            if "uddg" in qs:
+                url = unquote(qs["uddg"][0])
+        urls.append(url.rstrip(".,;"))
+    return urls
+
+
+def search(query: str, pages: int = 2) -> Iterable[str]:
+    endpoints = [
+        "https://www.bing.com/search?q={query}&first={offset}",
+        "https://html.duckduckgo.com/html/?q={query}&s={offset}",
+    ]
+    for endpoint in endpoints:
+        for page in range(pages):
+            offset = page * 10 + 1
+            url = endpoint.format(query=quote_plus(query), offset=offset)
+            try:
+                response = requests.get(url, headers=HEADERS, timeout=20)
+                if response.status_code >= 400:
+                    continue
+                yield from extract_urls(response.text)
+                time.sleep(0.4)
+            except requests.RequestException:
+                continue
+
+
+def github_json(url: str) -> dict | None:
+    try:
+        response = requests.get(url, headers={**HEADERS, "Accept": "application/vnd.github+json"}, timeout=25)
+        if response.status_code >= 400:
+            return None
+        return response.json()
+    except requests.RequestException:
+        return None
+
+
+def github_docx_urls(max_repos_per_query: int = 8) -> Iterable[str]:
+    seen_repos: set[str] = set()
+    for query in GITHUB_REPO_QUERIES:
+        search_url = (
+            "https://api.github.com/search/repositories"
+            f"?q={quote_plus(query)}&per_page={max_repos_per_query}"
+        )
+        payload = github_json(search_url)
+        if not payload:
+            continue
+        for repo in payload.get("items", []):
+            full_name = repo.get("full_name")
+            branch = repo.get("default_branch") or "main"
+            if not full_name or full_name in seen_repos:
+                continue
+            seen_repos.add(full_name)
+            tree_url = f"https://api.github.com/repos/{full_name}/git/trees/{branch}?recursive=1"
+            tree = github_json(tree_url)
+            if not tree:
+                continue
+            for item in tree.get("tree", []):
+                path = item.get("path", "")
+                if item.get("type") == "blob" and path.lower().endswith(".docx"):
+                    yield f"https://raw.githubusercontent.com/{full_name}/{branch}/{quote_path(path)}"
+            time.sleep(0.25)
+
+
+def quote_path(path: str) -> str:
+    return "/".join(quote_plus(part).replace("+", "%20") for part in path.split("/"))
+
+
+def looks_like_docx_url(url: str) -> bool:
+    low = url.lower()
+    return ".docx" in low and not any(bad in low for bad in ("?format=pdf", "/view?", "webcache"))
+
+
+def safe_name(index: int, url: str, content: bytes) -> str:
+    parsed = urlparse(url)
+    stem = Path(unquote(parsed.path)).name or f"online-{index:03d}.docx"
+    stem = re.sub(r"[^A-Za-z0-9._ -]+", "_", stem)
+    if not stem.lower().endswith(".docx"):
+        stem += ".docx"
+    digest = hashlib.sha256(content).hexdigest()[:10]
+    return f"{index:03d}-{digest}-{stem}"
+
+
+def validate_docx(path: Path) -> dict:
+    with ZipFile(path) as zf:
+        names = set(zf.namelist())
+        if "[Content_Types].xml" not in names or "word/document.xml" not in names:
+            raise BadZipFile("not a Word DOCX package")
+        document_xml = zf.read("word/document.xml")
+        return {
+            "omml_nodes": len(re.findall(rb"<m:oMath\b|<m:oMathPara\b", document_xml)),
+            "media_parts": len([name for name in names if name.startswith("word/media/")]),
+            "document_xml_bytes": len(document_xml),
+        }
+
+
+def download(url: str, out_dir: Path, index: int, max_mb: int) -> dict | None:
+    try:
+        with requests.get(url, headers=HEADERS, timeout=30, stream=True, allow_redirects=True) as response:
+            if response.status_code >= 400:
+                return None
+            content_type = response.headers.get("content-type", "").split(";")[0].strip().lower()
+            if content_type and content_type not in DOCX_CT_HINTS and "wordprocessingml" not in content_type:
+                if ".docx" not in response.url.lower():
+                    return None
+            chunks: list[bytes] = []
+            total = 0
+            limit = max_mb * 1024 * 1024
+            for chunk in response.iter_content(chunk_size=1024 * 256):
+                if not chunk:
+                    continue
+                total += len(chunk)
+                if total > limit:
+                    return None
+                chunks.append(chunk)
+        content = b"".join(chunks)
+    except requests.RequestException:
+        return None
+
+    if not content.startswith(b"PK"):
+        return None
+
+    out_path = out_dir / safe_name(index, response.url, content)
+    out_path.write_bytes(content)
+    try:
+        stats = validate_docx(out_path)
+    except (BadZipFile, KeyError):
+        out_path.unlink(missing_ok=True)
+        return None
+
+    return {
+        "file": str(out_path),
+        "source_url": response.url,
+        "bytes": len(content),
+        "content_type": content_type,
+        **stats,
+    }
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--out", type=Path, required=True)
+    parser.add_argument("--target", type=int, default=50)
+    parser.add_argument("--max-mb", type=int, default=40)
+    parser.add_argument("--prefer-equations", action="store_true")
+    args = parser.parse_args()
+
+    args.out.mkdir(parents=True, exist_ok=True)
+    seen: set[str] = set()
+    manifest: list[dict] = []
+    index = len(list(args.out.glob("*.docx"))) + 1
+
+    print("source: github repositories", flush=True)
+    for url in github_docx_urls():
+        if len(manifest) >= args.target:
+            break
+        normalized = url.split("#", 1)[0]
+        if normalized in seen:
+            continue
+        seen.add(normalized)
+        item = download(normalized, args.out, index, args.max_mb)
+        if not item:
+            continue
+        manifest.append(item)
+        index += 1
+        print(f"kept {len(manifest):02d}: {Path(item['file']).name} omml={item['omml_nodes']}", flush=True)
+
+    for query in DEFAULT_QUERIES:
+        print(f"search: {query}", flush=True)
+        for url in search(query):
+            if len(manifest) >= args.target:
+                break
+            if not looks_like_docx_url(url):
+                continue
+            normalized = url.split("#", 1)[0]
+            if normalized in seen:
+                continue
+            seen.add(normalized)
+            item = download(normalized, args.out, index, args.max_mb)
+            if not item:
+                continue
+            if args.prefer_equations and item["omml_nodes"] == 0 and len(manifest) < args.target // 2:
+                Path(item["file"]).unlink(missing_ok=True)
+                continue
+            manifest.append(item)
+            index += 1
+            print(f"kept {len(manifest):02d}: {Path(item['file']).name} omml={item['omml_nodes']}", flush=True)
+        if len(manifest) >= args.target:
+            break
+
+    manifest_path = args.out / "online-docx-manifest.json"
+    manifest_path.write_text(json.dumps({"count": len(manifest), "items": manifest}, indent=2), encoding="utf-8")
+    print(f"saved manifest: {manifest_path}", flush=True)
+    return 0 if len(manifest) >= args.target else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())