[Python] Topic modelling many documents with low memory overhead

Discussão em 'Python' iniciado por Stack, Fevereiro 20, 2025.

  1. Stack

    Stack Membro Participativo

    I've been working on a topic modelling project using BERTopic 0.16.3, and the preliminary results were promising. However, as the project progressed and the requirements became apparent, I ran into a specific issue with scalability.

    Specifically:

    • For development/testing, it needs to train reasonably quickly on a moderate number of documents (tens of thousands to low hundred thousands)
      • Our dev machines are Macs, so this probably has to be done on CPU
    • For production, it needs to train on a large number of documents (several million) without blowing up memory usage
      • For a baseline, with the default settings on my machine, BERTopic has a peak memory usage of roughly 35 kB per document, which easily becomes hundreds of GBs or even TBs for the amount of data that will be provided in production
      • Ideally, this would have peak memory usage sublinear in the number of documents.

    That last requirement necessitates batching the documents, since loading them all into memory at once requires linear memory. So, I've been looking into clustering algorithms that work with online topic modelling. BERTopic's documentation suggests scikit-learn's MiniBatchKMeans, but the results I'm getting from that aren't very good.

    Some models I've looked at include:

    • Birch via scikit-learn: uses even more memory than BERTopic's default HDBSCAN even when batched. Also runs much slower.
    • IncrementalDBSCAN via incdbscan: Seemed promising at first, but the runtime and eventually memory ballooned. For ~120k documents in batches of 5000, it didn't use more than 4GB of RAM in the first 3½ hours, but didn't finish within ten hours, and used nearly 40GB of RAM at some point in the middle.
    • AgglomerativeClustering via scikit-learn: gave very good results from initial testing (perhaps even better than HDBSCAN), but it doesn't implement the partial_fit method. I found this answer on a different question which suggests it's possible to train two of them using single linkage independently and then merge them, but it gives no indication as to how.

    The latter two also don't provide the predict method, limiting their utility.

    I am fairly new to the subject, so perhaps I'm approaching this completely wrong and the immediate problem I'm trying to solve has no solution. So to be clear, at the base level, the question I'm trying to answer is: How do I perform topic modelling (and get good results) on a large number of documents without using too much memory?

    Continue reading...

Compartilhe esta Página