Function rorf_classifier

trainer.py:38–87

(
    ctx,
    model_a: str,
    model_b: str,
    dataset_path: str,
    eval_dataset: str,
    embedding_provider: str,
    prompt_embedding_cache: Optional[str],
    max_depth: int,
    max_features: Union[float, str],
    n_estimators: int,
    save_dir: str,
    model_id: str,
    model_org: str,
)

@click.pass_context
def rorf_classifier(
    ctx,
    model_a: str,
    model_b: str,
    dataset_path: str,
    eval_dataset: str,
    embedding_provider: str,
    prompt_embedding_cache: Optional[str],
    max_depth: int,
    max_features: Union[float, str],
    n_estimators: int,
    save_dir: str,
    model_id: str,
    model_org: str,
):
    from rorf.router.rorf import RoRFTrainer

    llms = [model_a, model_b]

    configs = {
        "trainer": "RoRF",
        "llms": llms,
        "dataset_path": dataset_path,
        "eval_dataset": eval_dataset,
        "embedding_provider": embedding_provider,
        "prompt_embedding_cache": prompt_embedding_cache,
        "max_depth": max_depth,
        "max_features": max_features,
        "n_estimators": n_estimators,
        "save_dir": save_dir,
        "model_id": model_id,
        "model_org": model_org,
    }

    trainer_obj = RoRFTrainer(
        llms=llms,
        dataset_path=dataset_path,
        eval_dataset=eval_dataset,
        embedding_provider=embedding_provider,
        prompt_embedding_cache=prompt_embedding_cache,
        max_depth=max_depth,
        max_features=max_features,
        n_estimators=n_estimators,
        save_dir=save_dir,
        model_id=model_id,
        model_org=model_org,
    )
    write_config_to_json(configs, trainer_obj.save_path)
    ctx.obj["trainer"] = trainer_obj
    ctx.obj["configs"] = configs


@run.result_callback()