Function main

tests/modal_remote_vllm/remote_cache_dense.py:497–560

()

# Entry point
# ------------------------------------------------------------------
def main():
    # Declare globals at the top of the function
    global EMB_BATCH, CONCURRENT_BATCHES

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--dataset",
        choices=DATASETS,
        help="Run for a single dataset (optional). Default: all.",
    )
    parser.add_argument(
        "--use-chunking",
        action="store_true",
        help="Use text chunking instead of simple clipping for long documents",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=256,
        help="Batch size for embedding generation (default: 256)",
    )
    parser.add_argument(
        "--concurrent-batches",
        type=int,
        default=8,
        help="Number of concurrent API calls (default: 8)",
    )
    args = parser.parse_args()

    # Update globals based on arguments
    EMB_BATCH = args.batch_size
    CONCURRENT_BATCHES = args.concurrent_batches

    base_url = os.getenv(
        "REMOTE_BASE_URL", "https://rycerzes--arctic-embeddings-serve.modal.run/v1"
    )

    print(f"Using remote embedding server: {base_url}")
    print(
        f"Configuration: batch_size={EMB_BATCH}, concurrent_batches={CONCURRENT_BATCHES}"
    )

    # quick health-check
    try:
        print("Checking server health...")
        requests.get(base_url.replace("/v1", "/health"), timeout=5).raise_for_status()
        print("Server is healthy ✓")
    except Exception as e:
        parser.error(f"Remote server not reachable: {e}")

    print("Initializing tokenizer...")
    # Force tokenizer loading here to show progress
    _ = get_tokenizer().encode("test", add_special_tokens=True)
    print("Tokenizer ready ✓")

    embedder = RemoteEmbeddings(base_url, max_concurrent=CONCURRENT_BATCHES)

    todo = [args.dataset] if args.dataset else DATASETS