Add kmeans clustering based on ray #1080
Conversation
|
Don't we want the result of In the summarization case, we could use kmeans/clustering for topic discovery and select a few elements from each cluster and use those to summarize. What I would need is membership info. |
|
Technically, this is K means and I am talking about clustering. |
The github comment seems not work, your comment could not be replied in dialogue. Anyway, the PR already shows how the clustering method works. Basically, you do a similar trick to assign each row a cluster id using a map function by comparing with the centroids, and then do either sort or aggregate for your purpose. In your summarization case, you might chain an aggregate function to select N entries out for each group. |
lib/sycamore/sycamore/docset.py
Outdated
| @@ -903,6 +906,28 @@ def map(self, f: Callable[[Document], Document], **resource_args) -> "DocSet": | |||
| mapping = Map(self.plan, f=f, **resource_args) | |||
| return DocSet(self.context, mapping) | |||
|
|
|||
| def kmeans(self, K: int, iterations: int, init_mode: str = "random", epsilon: float = 1e-4): | |||
There was a problem hiding this comment.
I get that it's different, but I do think we should have a method that returns a Docset with an extra field indicating that cluster that each row is assigned to. I know that it is just an extra map, but I expect it's how people would want to use it in practice and it makes things more chainable.
There was a problem hiding this comment.
yes, sure, would add one.
|
|
||
| Args: | ||
| K: the count of centroids | ||
| iterations: the max iteration runs before converge |
There was a problem hiding this comment.
Is there a reasonable default for this? I at least wouldn't know what a good value to pick would be.
There was a problem hiding this comment.
spark uses 20, we could follow the same, but it should really be a tuning process.
|
How hard would it be to add a local mode version of this? I guess it would require a local-mode aggregate. That would be convenient, though I confess I'm not sure how important it is. |
yes, would fix. |
would try to come up with one. |
This includes generally three steps: 1. materialize a document's embedding 2. initialize centroids randomly 2. iterate the kmeans process until converge, this is based on ray dataset map group and aggregate operators. The result centroids could be used for downstream work.
This includes generally three steps:
The result centroids could be used for downstream work.