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

Code before:
    )
    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:
        if args.model == "text-embedding-ada-002":
            texts = read_texts(
                args.input_file, args.min_length, args.max_length, args.english_only
            )
            print(f"#text: {len(texts)}")
            embeddings = get_openai_embeddings(texts, args.model, args.batch_size)
            print(f"embeddings shape: {embeddings.shape}")
        else:
            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)
    with open(filename_prefix + "_topk.txt", "w") as fout:
        fout.write(topk_str)

Review comment (performance):
```suggestion
        texts = read_texts(
            args.input_file, args.min_length, args.max_length, args.english_only
        )
        print(f"#text: {len(texts)}")
        if args.model == "text-embedding-ada-002":
            embeddings = get_openai_embeddings(texts, args.model, args.batch_size)
        else:
            embeddings = get_embeddings(texts, args.model, args.batch_size)
        print(f"embeddings shape: {embeddings.shape}")
```

Code after:
    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}"
    print(topk_str)
    with open(filename_prefix + "_topk.txt", "w") as fout:
        fout.write(topk_str)

    with open(filename_prefix + "_all.jsonl", "w") as fout:
        for i in range(len(centers)):
            tmp_indices = labels == i
            tmp_embeddings = embeddings[tmp_indices]
            tmp_texts = texts[tmp_indices]

            scores = cos_sim(centers[i].unsqueeze(0), tmp_embeddings)[0]
            sorted_indices = torch.flip(torch.argsort(scores), dims=[0])

            for text, score in zip(tmp_texts[sorted_indices], scores[sorted_indices]):