Repository: lm-sys/FastChat
File: fastchat/serve/monitor/topic_clustering.py (Python)

Code before:
    # default="multi-qa-distilbert-cos-v1")
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--min-length", type=int)
    parser.add_argument("--max-length", type=int)
    parser.add_argument("--english-only", action="store_true")
    parser.add_argument("--num-clusters", type=int, default=20)
    parser.add_argument(
        "--cluster-alg",
        type=str,
        choices=["kmeans", "aggcls", "HDBSCAN"],
        default="kmeans",
    )
    parser.add_argument("--show-top-k", type=int, default=200)
    parser.add_argument("--show-cut-off", type=int, default=512)
    parser.add_argument("--save-embeddings", action="store_true")
    parser.add_argument("--embeddings-file", type=str, default=None)
    args = parser.parse_args()

    num_clusters = args.num_clusters
    show_top_k = args.show_top_k
    show_cut_off = args.show_cut_off

    if args.embeddings_file is None:
        texts = read_texts(
            args.input_file, args.min_length, args.max_length, args.english_only
        )
        print(f"#text: {len(texts)}")
        embeddings = get_embeddings(texts, args.model, args.batch_size)
        print(f"embeddings shape: {embeddings.shape}")
        if args.save_embeddings:
            # allow saving embedding to save time
            torch.save(embeddings, "embeddings.pt")
    else:
        embeddings = torch.load(args.embeddings_file)

    if args.cluster_alg == "kmeans":
        centers, labels = run_k_means(embeddings, num_clusters)
    elif args.cluster_alg == "aggcls":
        centers, labels = run_agg_cluster(embeddings, num_clusters)
    elif args.cluster_alg == "HDBSCAN":
        centers, labels = run_hdbscan_cluster(embeddings)
    else:
        raise ValueError(f"Invalid clustering algorithm: {args.cluster_alg}")

    topk_indices = get_topk_indices(centers, labels, embeddings, args.show_top_k)
    topk_str = print_topk(texts, labels, topk_indices, args.show_cut_off)
    num_clusters = len(centers)

    # Dump results
    filename_prefix = f"results_c{num_clusters}_{args.cluster_alg}"
    print(topk_str)

Review comment (suggestion):
This needs to be move out from `if`. otherwise if you load embedding from file, `texts` won't be correctly initialized

Code after:
    # default="multi-qa-distilbert-cos-v1")
    parser.add_argument("--batch-size", type=int, default=256)
    parser.add_argument("--min-length", type=int)
    parser.add_argument("--max-length", type=int)
    parser.add_argument("--english-only", action="store_true")
    parser.add_argument("--num-clusters", type=int, default=20)
    parser.add_argument(
        "--cluster-alg",
        type=str,
        choices=["kmeans", "aggcls", "HDBSCAN"],
        default="kmeans",
    )
    parser.add_argument("--show-top-k", type=int, default=200)
    parser.add_argument("--show-cut-off", type=int, default=512)
    parser.add_argument("--save-embeddings", action="store_true")
    parser.add_argument("--embeddings-file", type=str, default=None)
    args = parser.parse_args()

    num_clusters = args.num_clusters
    show_top_k = args.show_top_k
    show_cut_off = args.show_cut_off

    texts = read_texts(
        args.input_file, args.min_length, args.max_length, args.english_only
    )
    print(f"#text: {len(texts)}")

    if args.embeddings_file is None:
        embeddings = get_embeddings(texts, args.model, args.batch_size)
        if args.save_embeddings:
            # allow saving embedding to save time and money
            torch.save(embeddings, "embeddings.pt")
    else:
        embeddings = torch.load(args.embeddings_file)
    print(f"embeddings shape: {embeddings.shape}")

    if args.cluster_alg == "kmeans":
        centers, labels = run_k_means(embeddings, num_clusters)
    elif args.cluster_alg == "aggcls":
        centers, labels = run_agg_cluster(embeddings, num_clusters)
    elif args.cluster_alg == "HDBSCAN":
        centers, labels = run_hdbscan_cluster(embeddings)
    else:
        raise ValueError(f"Invalid clustering algorithm: {args.cluster_alg}")

    topk_indices = get_topk_indices(centers, labels, embeddings, args.show_top_k)
    topk_str = print_topk(texts, labels, topk_indices, args.show_cut_off)
    num_clusters = len(centers)

    # Dump results
    filename_prefix = f"results_c{num_clusters}_{args.cluster_alg}"