run_proxy_sft.py

import typer
from pathlib import Path

from sse_bio.training.proxy_data import export_proxy_training_rows
from sse_bio.training.sft import train_proxy_sft

app = typer.Typer(add_completion=False, help="Build SSE-Bio proxy SFT data and train a proxy backbone.")


@app.command("build-data")
def build_data(
    dataset_path: Path = typer.Argument(..., help="Path to a BioHopR JSON/JSONL file or bundle directory."),
    output_path: Path = typer.Option(..., "--output-path", help="Path to save proxy SFT jsonl."),
    split: str = typer.Option("train", "--split", help="BioHopR split used for proxy training data."),
):
    count = export_proxy_training_rows(dataset_path, output_path, split=split, mode="sft")
    typer.echo(f"Exported {count} proxy SFT rows to {output_path}")


@app.command("train")
def train(
    dataset_path: Path = typer.Argument(..., help="Path to proxy SFT jsonl."),
    model_name_or_path: str = typer.Option(..., "--model", help="Proxy backbone, e.g. Qwen/Qwen2.5-7B-Instruct."),
    output_dir: Path = typer.Option(..., "--output-dir", help="Directory for trained SFT checkpoint."),
    max_seq_length: int = typer.Option(2048, "--max-seq-length"),
    learning_rate: float = typer.Option(5e-6, "--learning-rate"),
    num_train_epochs: int = typer.Option(1, "--epochs"),
    per_device_train_batch_size: int = typer.Option(1, "--batch-size"),
    gradient_accumulation_steps: int = typer.Option(0, "--grad-accum", help="Auto-computed when set to 0."),
    weight_decay: float = typer.Option(0.01, "--weight-decay"),
    warmup_ratio: float = typer.Option(0.03, "--warmup-ratio"),
    lr_scheduler_type: str = typer.Option("cosine", "--lr-scheduler"),
    target_effective_batch_size: int = typer.Option(64, "--target-effective-batch-size"),
    save_steps: int = typer.Option(400, "--save-steps"),
    logging_steps: int = typer.Option(10, "--logging-steps"),
):
    train_proxy_sft(
        dataset_path=dataset_path,
        model_name_or_path=model_name_or_path,
        output_dir=output_dir,
        max_seq_length=max_seq_length,
        learning_rate=learning_rate,
        num_train_epochs=num_train_epochs,
        per_device_train_batch_size=per_device_train_batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        weight_decay=weight_decay,
        warmup_ratio=warmup_ratio,
        lr_scheduler_type=lr_scheduler_type,
        target_effective_batch_size=target_effective_batch_size,
        save_steps=save_steps,
        logging_steps=logging_steps,
    )


if __name__ == "__main__":
    app()