SPB Git forge

spb/ai-atlas

Public
41commits 1branches 0releases
4.6 MBsize
maindefault branch
12 days agolast push
HTML 77.2% TypeScript 10.5% Python 9.6% JavaScript 2.5%
3.5 KB · 102 lines python
Raw Blame History
1"""Markdown documents (model cards, docs served as .md, READMEs): YAML front matter, headings, tables, links, plain text."""2from __future__ import annotations34import re5from dataclasses import dataclass, field6from typing import Any78import yaml910_FM = re.compile(r"\A---[ \t]*\r?\n(.*?)\r?\n---[ \t]*\r?\n", re.DOTALL)11_HEADING = re.compile(r"^(#{1,6})\s+(.*?)\s*#*\s*$", re.MULTILINE)12_LINK = re.compile(r"\[([^\]]*)\]\((https?://[^)\s]+)\)")13_TABLE_ROW = re.compile(r"^\s*\|(.+)\|\s*$")14_SEP_ROW = re.compile(r"^\s*\|?\s*:?-{2,}:?\s*(\|\s*:?-{2,}:?\s*)*\|?\s*$")151617@dataclass18class MarkdownDoc:19    front_matter: dict[str, Any] = field(default_factory=dict)20    body: str = ""21    headings: list[tuple[int, str]] = field(default_factory=list)22    tables: list[dict[str, Any]] = field(default_factory=list)23    links: list[tuple[str, str]] = field(default_factory=list)24    text: str = ""2526    def section(self, title_pattern: str) -> str:27        """Body text of the first heading matching `title_pattern` (case-insensitive) up to the next heading of same/higher level."""28        rx = re.compile(title_pattern, re.IGNORECASE)29        lines = self.body.split("\n")30        start = None31        level = 032        for i, line in enumerate(lines):33            m = _HEADING.match(line)34            if m and start is None and rx.search(m.group(2)):35                start = i + 136                level = len(m.group(1))37                continue38            if start is not None and m and len(m.group(1)) <= level:39                return "\n".join(lines[start:i]).strip()40        return "\n".join(lines[start:]).strip() if start is not None else ""414243def parse_front_matter(text: str) -> tuple[dict[str, Any], str]:44    m = _FM.match(text)45    if not m:46        return {}, text47    try:48        data = yaml.safe_load(m.group(1)) or {}49        if not isinstance(data, dict):50            data = {"_value": data}51    except yaml.YAMLError:52        data = {}53    return data, text[m.end():]545556def parse_markdown(text: str) -> MarkdownDoc:57    fm, body = parse_front_matter(text)58    doc = MarkdownDoc(front_matter=fm, body=body)59    doc.headings = [(len(m.group(1)), m.group(2).strip()) for m in _HEADING.finditer(body)]60    doc.links = [(m.group(2), m.group(1)) for m in _LINK.finditer(body)]61    doc.tables = _tables(body)62    doc.text = _to_text(body)63    return doc646566def _tables(body: str) -> list[dict[str, Any]]:67    out: list[dict[str, Any]] = []68    lines = body.split("\n")69    i = 070    while i < len(lines):71        if _TABLE_ROW.match(lines[i]) and i + 1 < len(lines) and _SEP_ROW.match(lines[i + 1]):72            headers = _cells(lines[i])73            rows: list[list[str]] = []74            i += 275            while i < len(lines) and _TABLE_ROW.match(lines[i]):76                rows.append(_cells(lines[i]))77                i += 178            out.append({"caption": None, "headers": headers, "rows": rows})79        else:80            i += 181    return out828384def _cells(line: str) -> list[str]:85    inner = line.strip()86    inner = inner.removeprefix("|")87    inner = inner.removesuffix("|")88    return [re.sub(r"\*\*|`", "", c).strip() for c in inner.split("|")]899091def _to_text(body: str) -> str:92    s = re.sub(r"```.*?```", " ", body, flags=re.DOTALL)93    s = re.sub(r"!\[[^\]]*\]\([^)]*\)", " ", s)94    s = re.sub(r"\[([^\]]*)\]\([^)]*\)", r"\1", s)95    s = re.sub(r"<[^>]+>", " ", s)96    s = re.sub(r"^[#>*\-|\s]+", "", s, flags=re.MULTILINE)97    s = re.sub(r"[ \t]+", " ", s)98    return re.sub(r"\n{3,}", "\n\n", s).strip()99100101__all__ = ["MarkdownDoc", "parse_front_matter", "parse_markdown"]102