marin-core 0.99__py3-none-any.whl

marin/__init__.py

@@ -0,0 +1,2 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0

marin/cluster/__init__.py

@@ -0,0 +1,10 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Unified cluster management utilities."""
+
+from . import gcp
+
+__all__ = [
+    "gcp",
+]

marin/cluster/gcp.py

@@ -0,0 +1,158 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""GCP utilities for cluster management.
+
+This provides functions to get access to the current GCP configuration, list and
+connect to TPUs, and find TPUs by IP address.
+"""
+
+import json
+import logging
+import subprocess
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def run_gcloud_command(cmd: list[str], **kwargs) -> subprocess.CompletedProcess:
+    """Run a gcloud command with error handling."""
+    try:
+        logger.info(f"Running {' '.join(cmd)}")
+        return subprocess.run(cmd, check=True, capture_output=True, text=True, **kwargs)
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"gcloud command failed: {' '.join(cmd)}\nError: {e.stderr}") from e
+
+
+def get_project_id() -> str | None:
+    """Get the current GCP project ID."""
+    try:
+        result = run_gcloud_command(["gcloud", "config", "get-value", "project"])
+        return result.stdout.strip() or None
+    except RuntimeError:
+        return None
+
+
+def list_tpu_nodes(project: str, zone: str, filter_expr: str = "") -> list[dict[str, Any]]:
+    """List TPU nodes in a zone."""
+    cmd = [
+        "gcloud",
+        "compute",
+        "tpus",
+        "tpu-vm",
+        "list",
+        f"--project={project}",
+        f"--zone={zone}",
+        "--format=json",
+    ]
+    if filter_expr:
+        cmd.append(f"--filter={filter_expr}")
+
+    result = run_gcloud_command(cmd)
+    return json.loads(result.stdout)
+
+
+def find_tpu_by_ip(target_ip: str, project: str, zone: str = "-") -> tuple[str, str, int] | None:
+    """Find TPU node by its internal IP address.
+
+    Searches all zones by default (zone="-").
+
+    Returns:
+        Tuple of (tpu_name, zone, worker_index) or None if not found
+    """
+    tpu_nodes = list_tpu_nodes(project, zone)
+
+    for node in tpu_nodes:
+        network_endpoints = node.get("networkEndpoints", [])
+        for worker_index, endpoint in enumerate(network_endpoints):
+            if endpoint.get("ipAddress") == target_ip:
+                # Extract simple name from full resource path
+                full_name = node["name"]
+                name_parts = full_name.split("/")
+                if len(name_parts) >= 6:
+                    simple_name = name_parts[5]  # nodes/simple-name
+                    node_zone = name_parts[3]  # locations/zone
+                    return simple_name, node_zone, worker_index
+                else:
+                    # Fallback for different naming schemes
+                    return full_name, zone, worker_index
+
+    return None
+
+
+def find_vm_by_ip(target_ip: str, project: str) -> tuple[str, str] | None:
+    """Find a GCE VM by its internal IP address.
+
+    Searches all zones.
+
+    Returns:
+        Tuple of (instance_name, zone) or None if not found
+    """
+    cmd = [
+        "gcloud",
+        "compute",
+        "instances",
+        "list",
+        f"--project={project}",
+        f"--filter=networkInterfaces[0].networkIP={target_ip}",
+        "--format=json(name,zone)",
+    ]
+    result = run_gcloud_command(cmd)
+    instances = json.loads(result.stdout)
+    if not instances:
+        return None
+
+    if len(instances) > 1:
+        details = ", ".join(f"{i['name']} ({i['zone'].split('/')[-1]})" for i in instances)
+        raise RuntimeError(f"Multiple VMs found with IP {target_ip}: {details}")
+
+    instance = instances[0]
+    name = instance["name"]
+    # zone is a full URL like .../zones/us-central1-a
+    zone = instance["zone"].split("/")[-1]
+    return name, zone
+
+
+def ssh_to_vm(instance_name: str, zone: str, project: str, extra_args: list[str] | None = None) -> None:
+    """SSH into a GCE VM."""
+    cmd = [
+        "gcloud",
+        "compute",
+        "ssh",
+        instance_name,
+        f"--zone={zone}",
+        f"--project={project}",
+    ]
+
+    if extra_args:
+        cmd.extend(["--", *extra_args])
+
+    subprocess.run(cmd, check=True)
+
+
+def ssh_to_tpu(tpu_name: str, zone: str, project: str, extra_args: list[str] | None = None, worker_id: int = 0) -> None:
+    """SSH into a TPU node.
+
+    Args:
+        tpu_name: Name of the TPU
+        zone: GCP zone
+        project: GCP project
+        extra_args: Additional SSH arguments
+        worker_id: Worker index for multi-worker TPUs (default: 0)
+    """
+    cmd = [
+        "gcloud",
+        "compute",
+        "tpus",
+        "tpu-vm",
+        "ssh",
+        tpu_name,
+        f"--zone={zone}",
+        f"--project={project}",
+        f"--worker={worker_id}",
+    ]
+
+    if extra_args:
+        cmd.extend(["--", *extra_args])
+
+    subprocess.run(cmd, check=True)

marin/core/conversation.py

@@ -0,0 +1,25 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class OpenAIChatMessage(BaseModel):
+    model_config = ConfigDict(extra="allow")
+    role: str
+    content: Any
+    name: str | None = None
+    tool_calls: list[dict[str, Any]] | None = None
+    tool_call_id: str | None = Field(default=None, alias="tool_call_id")
+
+
+class DolmaConversationOutput(BaseModel):
+    model_config = ConfigDict(extra="allow")
+    id: str
+    source: str
+    messages: list[OpenAIChatMessage]
+    added: str
+    created: str
+    metadata: dict[str, Any]

marin/core/data.py

@@ -0,0 +1,54 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass
+
+
+@dataclass
+class QAExampleMetadata:
+    """
+    Dataclass representing metadata for a document.
+
+    Attributes:
+        subset (str): subset of dataset
+        split (str): split of dataset
+        revision (str): revision of dataset
+        provenance (str): URL of source of data, usually HF
+        answer (str): text of answer
+        answer_idx (str): index into list of answer options corresponding to correct answer
+        answer_label (str): label of correct answer (e.g. A)
+        options (list[str]): list of potential options for multiple choice question
+        answer_labels (list[str]): list of labels for options (e.g. A,B,C,D)
+    """
+
+    subset: str | None = None
+    split: str | None = None
+    revision: str | None = None
+    provenance: str | None = None
+    answer: str | None = None
+    answer_idx: int | None = None
+    answer_label: str | None = None
+    options: list[str] | None = None
+    answer_labels: list[str] | None = None
+
+
+@dataclass
+class QAExample:
+    """
+    Dataclass representing a document
+
+    Attributes:
+        id (str): Unique identifier for the record.
+        source (str): The name of the dataset.
+        metadata (QAExampleMetadata): Metadata related to the dataset.
+        text (str): The text of the document
+        prompt (str): If document is prompt/response, the prompt component
+        response (str): If document is prompt/response, the expected response component
+    """
+
+    id: str
+    source: str
+    metadata: QAExampleMetadata
+    text: str | None = None
+    prompt: str | None = None
+    response: str | None = None

marin/datakit/__init__.py

@@ -0,0 +1,24 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Datakit: composable pipeline stages with a standard Parquet format.
+
+The standard format pins three mandatory columns on every normalized record:
+``id`` (deterministic content hash), ``text`` (UTF-8 primary content), and
+``partition_id`` (int, the output shard the row was written to at normalize
+time). The shard count itself lives on the artifact, not the row.
+
+Downstream stages preserve ``partition_id`` and use it as the ``group_by`` key
+when a global shuffle (e.g. cross-document dedup) needs to land output back
+co-partitioned with the source.
+"""
+
+
+def partition_filename(partition_id: int, num_partitions: int) -> str:
+    """Return the standard datakit partition filename for the given index.
+
+    Datakit shards follow ``part-NNNNN-of-MMMMM.parquet`` naming. Routing
+    output through this helper keeps shuffler-written attribute files
+    discoverable by consolidate's filename-based join.
+    """
+    return f"part-{partition_id:05d}-of-{num_partitions:05d}.parquet"

marin/datakit/canonical/__init__.py

@@ -0,0 +1,2 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0

marin/datakit/canonical/fineweb_edu.py

@@ -0,0 +1,39 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Datakit canonical pipeline for FineWeb-Edu.
+
+HuggingFace: HuggingFaceFW/fineweb-edu
+
+FineWeb-Edu is a filtered subset of FineWeb selected for educational content.
+The raw download is Parquet with columns: text, id, url, dump, file_path,
+language, language_score, token_count, score, int_score.
+
+Subsets available on HuggingFace:
+- data/          — full dataset
+- sample/10BT    — 10B token sample
+- sample/100BT   — 100B token sample
+- sample/350BT   — 350B token sample
+"""
+
+from fray import ResourceConfig
+
+from marin.datakit.download.huggingface import download_hf_step
+from marin.execution.step_spec import StepSpec
+
+
+def download(
+    *,
+    revision: str = "87f0914",
+    hf_urls_glob: list[str] | None = None,
+    worker_resources: ResourceConfig | None = None,
+) -> StepSpec:
+    """Download FineWeb-Edu from HuggingFace."""
+    return download_hf_step(
+        "raw/fineweb-edu",
+        hf_dataset_id="HuggingFaceFW/fineweb-edu",
+        revision=revision,
+        hf_urls_glob=hf_urls_glob,
+        override_output_path=f"raw/fineweb-edu-{revision}",
+        worker_resources=worker_resources,
+    )

marin/datakit/download/__init__.py

@@ -0,0 +1,2 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0

marin/datakit/download/ar5iv.py

@@ -0,0 +1,156 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Download and process Ar5iv dataset from a zip file.
+
+"""
+
+import json
+import logging
+import zipfile
+from collections import defaultdict
+from dataclasses import dataclass
+
+import draccus
+from rigging.filesystem import open_url
+from rigging.log_setup import configure_logging
+from zephyr import Dataset, ZephyrContext
+from zephyr.writers import atomic_rename
+
+from marin.execution.step_spec import StepSpec
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Ar5ivDownloadConfig:
+    input_path: str
+    output_path: str
+    max_files: int | None = None  # Maximum number of shards to process
+
+
+def process_shard(shard_task: dict) -> dict:
+    """
+    Process a single shard by extracting its files from the zip in GCS and uploading the merged JSONL.
+
+    Args:
+        shard_task: Dict with keys 'input_path', 'output_path', 'shard_id', 'file_list'
+    """
+    input_path = shard_task["input_path"]
+    output_path = shard_task["output_path"]
+    shard_id = shard_task["shard_id"]
+    file_list = shard_task["file_list"]
+    gcs_path = f"{output_path}/{shard_id}.jsonl.gz"
+
+    with open_url(str(input_path), "rb") as f:
+        with zipfile.ZipFile(f) as zf:
+            with atomic_rename(gcs_path) as temp_path, open_url(temp_path, "wt", compression="gzip") as out_f:
+                for filename in file_list:
+                    with zf.open(filename, "r") as file_handle:
+                        content = file_handle.read()
+                        record = {
+                            "filename": filename,
+                            "format": "html",
+                            "content": content.decode("utf-8", errors="replace"),
+                        }
+                        print(json.dumps(record), file=out_f)
+
+            logger.info(f"Shard {shard_id} with {len(file_list)} files uploaded to {gcs_path}")
+            return {"shard_id": shard_id, "num_files": len(file_list), "output_path": gcs_path}
+
+
+def download(cfg: Ar5ivDownloadConfig) -> None:
+    """
+    Download and process Ar5iv dataset from a zip file in GCS.
+
+    This function can be called by the executor framework or used standalone.
+    """
+    logger.info("Starting transfer of Ar5iv dataset...")
+    logger.info(f"Source: {cfg.input_path}")
+
+    # Use fsspec+zipfile to list all files
+    with open_url(str(cfg.input_path), "rb") as f:
+        with zipfile.ZipFile(f) as zf:
+            all_files = zf.infolist()
+
+            # Group by shard directory
+            # We assume structure: something like: shard_id/.../file
+            # shard_id is derived from the second last component if files are nested.
+            # Adjust as needed if directory structure differs.
+            shard_dict = defaultdict(list)
+            for info in all_files:
+                if info.is_dir():
+                    continue
+                # E.g. path might look like: "003/something.html"
+                # Extract shard_id from the directory:
+                # Split by "/" and take the first part if we assume structure {shard_id}/file
+                parts = info.filename.strip("/").split("/")
+                if len(parts) < 2:
+                    # File at root level - decide how to handle this case.
+                    # If no directory structure is given, skip or treat differently.
+                    continue
+                shard_id = parts[-2]  # get the second-last directory as shard_id
+                shard_dict[shard_id].append(info.filename)
+
+            # Apply max_files limit if provided
+            shard_ids = list(shard_dict.keys())
+            if cfg.max_files is not None:
+                shard_ids = shard_ids[: cfg.max_files]
+
+            logger.info(f"Found {len(shard_ids)} shards to process.")
+
+            # Build task list for each shard
+            shard_tasks = []
+            for shard_id in shard_ids:
+                shard_tasks.append(
+                    {
+                        "input_path": cfg.input_path,
+                        "output_path": cfg.output_path,
+                        "shard_id": shard_id,
+                        "file_list": shard_dict[shard_id],
+                    }
+                )
+
+    # Execute pipeline with zephyr
+    pipeline = (
+        Dataset.from_list(shard_tasks)
+        .map(process_shard)
+        .write_jsonl(f"{cfg.output_path}/.metrics/part-{{shard:05d}}.jsonl", skip_existing=True)
+    )
+    ctx = ZephyrContext(name="download-ar5iv")
+    ctx.execute(pipeline)
+
+    logger.info("Transfer completed successfully!")
+
+
+def ar5iv_step(
+    name: str = "raw/ar5iv",
+    *,
+    input_path: str,
+    max_files: int | None = None,
+    deps: list[StepSpec] | None = None,
+    output_path_prefix: str | None = None,
+    override_output_path: str | None = None,
+) -> StepSpec:
+    """Create a StepSpec that downloads and processes the Ar5iv dataset from a zip file."""
+
+    def _run(output_path: str) -> None:
+        download(Ar5ivDownloadConfig(input_path=input_path, output_path=output_path, max_files=max_files))
+
+    return StepSpec(
+        name=name,
+        fn=_run,
+        deps=deps or [],
+        hash_attrs={"input_path": input_path, "max_files": max_files},
+        output_path_prefix=output_path_prefix,
+        override_output_path=override_output_path,
+    )
+
+
+@draccus.wrap()
+def main(cfg: Ar5ivDownloadConfig) -> None:
+    """CLI entrypoint for downloading and processing Ar5iv dataset."""
+
+    configure_logging(level=logging.INFO)
+    download(cfg)

marin/datakit/download/bio_chem/__init__.py

@@ -0,0 +1,22 @@
+# Copyright The Marin Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Streaming downloaders for biology and chemistry notation slices.
+
+Each submodule defines a StepSpec factory for one source family
+(RefSeq, RNAcentral, UniProt, PubChem, RCSB PDB, ChEMBL, MoleculeNet) that
+streams from the upstream mirror, splits the stream into format-preserving
+records via :mod:`marin.transform.bio_chem`, packs short records into longer
+documents for in-context-learning evaluation, and writes the result to
+plain-text-in-parquet that Levanter can read directly.
+
+The shared streaming primitives live in :mod:`._runtime`.
+"""
+
+from marin.datakit.download.bio_chem._runtime import (
+    NotationFormat,
+    NotationSliceSpec,
+    PackingConfig,
+    bio_chem_slice_step,
+    run_notation_slice,
+)