run_primekg.py

from pathlib import Path
from typing import List, Optional

import typer

from sse_bio.dense_retriever import export_dense_embeddings
from sse_bio.primekg import default_primekg_allowed_types, export_primekg_triplets
from sse_bio.triplet_store import TripletStore

app = typer.Typer(add_completion=False, help="PrimeKG preprocessing utilities for SSE-Bio.")


@app.callback()
def main() -> None:
    """Expose PrimeKG utilities in subcommand mode."""


@app.command("convert")
def convert(
    primekg_csv_path: Path = typer.Argument(..., help="Path to PrimeKG kg.csv."),
    output_path: Path = typer.Option(..., "--output-path", help="Output JSONL triplet file."),
    allowed_types: Optional[List[str]] = typer.Option(
        None,
        "--allowed-type",
        help="Optional node-type filter. Repeat the option to include multiple PrimeKG node types.",
    ),
    use_default_allowed_types: bool = typer.Option(
        True,
        "--use-default-allowed-types/--all-types",
        help="Restrict to BioHopR-relevant node types by default.",
    ),
    limit: Optional[int] = typer.Option(None, "--limit", help="Optional row limit for quick inspection."),
):
    selected_types = allowed_types
    if not selected_types and use_default_allowed_types:
        selected_types = default_primekg_allowed_types()
    count = export_primekg_triplets(
        primekg_csv_path,
        output_path,
        allowed_types=selected_types,
        limit=limit,
    )
    typer.echo(f"Exported {count} SSE-Bio triplets to {output_path}")


@app.command("embed")
def embed(
    triplets_path: Path = typer.Argument(..., help="Path to a SSE-Bio triplet JSON/JSONL/CSV file."),
    output_path: Path = typer.Option(..., "--output-path", help="Output .pt embedding file."),
    encoder_name_or_path: str = typer.Option(..., "--encoder", help="Dense retriever encoder path or HF model id."),
    tokenizer_name_or_path: Optional[str] = typer.Option(None, "--tokenizer", help="Optional tokenizer path."),
    device: str = typer.Option("auto", "--device", help="Encoding device."),
    max_length: int = typer.Option(256, "--max-length"),
    batch_size: int = typer.Option(64, "--batch-size"),
    trust_remote_code: bool = typer.Option(False, "--trust-remote-code"),
):
    store = TripletStore(triplets_path)
    count = export_dense_embeddings(
        [triplet.search_text() for triplet in store.triplets],
        output_path,
        backend="bica_dense",
        encoder_name_or_path=encoder_name_or_path,
        tokenizer_name_or_path=tokenizer_name_or_path,
        device=device,
        trust_remote_code=trust_remote_code,
        max_length=max_length,
        batch_size=batch_size,
    )
    typer.echo(f"Exported dense embeddings for {count} SSE-Bio triplets to {output_path}")


if __name__ == "__main__":
    app()