[NewOp] Add domain_diversity_selector based on DaaR principles

[lingzhq]

Introduces a novel data selection op based on semantic diversity across domains, designed to automatically select the most diverse subset of data samples, which is inspired by the DaaR paper.

• Converts input samples into embeddings
• Use embeddings to cluster pseudo-domains
• Selects samples based on various distances to maximize diversity

[WIP] Ongoing development of additional operators derived from the DaaR.

[yxdyc]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a new domain_diversity_selector operator, which is a valuable addition for curating datasets based on semantic diversity. The implementation follows the DaaR paper's principles, including embedding generation, clustering, and strategic sample selection. The code is well-structured, and the inclusion of configuration, documentation, and tests is appreciated.

However, there are several critical areas for improvement, primarily concerning performance and memory efficiency. The current implementation processes samples one by one, which will be very slow on large datasets. I've provided suggestions to use batching and vectorized operations to significantly speed up embedding and similarity calculations. Additionally, the memory usage can be optimized by avoiding the storage of redundant data.

The tests are a good start, but they rely on external APIs, making them difficult to run in automated environments. I've suggested mocking these dependencies. Finally, a few minor refactoring opportunities and configuration improvements are noted to enhance code clarity and usability.

[gemini-code-assist]

The current implementation for generating embeddings processes samples one by one, which is highly inefficient for large datasets. Both HuggingFace encode methods and many embedding APIs support batch processing. Refactoring this to process samples in batches will significantly improve performance.

Suggested change

	if self.is_hf_model:
	# Embeddings extract via local models
	for sample in tqdm(dataset, desc="Embedding", unit="sample"):
	text = sample["text"]
	with torch.no_grad():
	embedding = model.encode(text)
	embeddings.append(embedding)
	else:
	# Embeddings extract via API
	for sample in tqdm(dataset, desc="Embedding", unit="sample"):
	text = sample["text"]
	embedding = model(text, dimensions=self.ebd_dim, encoding_format="float")
	embeddings.append(embedding)
	if self.is_hf_model:
	# Embeddings extract via local models in batches for efficiency
	texts = [sample['text'] for sample in dataset]
	with torch.no_grad():
	embeddings = model.encode(texts, batch_size=self.batch_size, show_progress_bar=True)
	else:
	# Embeddings extract via API. Consider batching if the API supports it.
	for sample in tqdm(dataset, desc="Embedding", unit="sample"):
	text = sample["text"]
	embedding = model(text, dimensions=self.ebd_dim, encoding_format="float")
	embeddings.append(embedding)

[gemini-code-assist]

This section has critical performance and memory issues that should be addressed:

1. Performance: Cosine similarities are calculated in a loop, creating and destroying PyTorch tensors for each sample. This is very slow. These calculations should be vectorized using sklearn.metrics.pairwise.cosine_similarity or torch.nn.functional.cosine_similarity on the entire matrices.
2. Memory: The data_status list stores a dictionary for each sample, including the original text. For large datasets, this will consume a very large amount of memory. It's more memory-efficient to work with NumPy arrays for embeddings, labels, and calculated similarities, and only use indices to refer back to the original dataset.

A full refactor is recommended to process these calculations in a batched/vectorized manner and to avoid creating large intermediate data structures.

[gemini-code-assist]

This test is skipped when running from a fork, and its implementation relies on an external API (text-embedding-v3), which requires credentials. This makes the test difficult to run in CI/CD environments and for external contributors. It's highly recommended to either mock the API calls or use a small, publicly available HuggingFace model (by setting is_hf_model=True and providing a model name) for testing to ensure the operator's logic can be verified automatically and reliably.

[gemini-code-assist]

The select_ratio is left empty in this configuration example. While the code correctly handles a None value by skipping the operation, it would be more user-friendly for an example configuration to provide a default value (e.g., 0.5) and add a comment explaining its purpose.

      select_ratio: 0.5                                           # the ratio to be sampled

[gemini-code-assist]

The if/elif/else chain for selecting the sorting key based on the strategy can be made more concise and extensible by using a dictionary to map strategy names to their corresponding sorting logic (key and reverse flag).

Suggested change

	if self.strategy == "inter":
	label_data_status.sort(key=lambda x: x["total_similarity"])
	elif self.strategy == "intra":
	label_data_status.sort(key=lambda x: x["similarity_with_own_centroid"], reverse=True)
	elif self.strategy == "global":
	label_data_status.sort(key=lambda x: x["global_centroid_similarity"])
	else:
	raise ValueError("Invalid strategy. Use 'inter', 'intra' or 'global'.")
	strategy_configs = {
	'inter': {'key': 'total_similarity', 'reverse': False},
	'intra': {'key': 'similarity_with_own_centroid', 'reverse': True},
	'global': {'key': 'global_centroid_similarity', 'reverse': False},
	}
	if self.strategy not in strategy_configs:
	raise ValueError("Invalid strategy. Use 'inter', 'intra' or 'global'.")
	config = strategy_configs[self.strategy]
	label_data_status.sort(key=lambda x: x[config['key']], reverse=config['reverse'])

[gemini-code-assist]

The assertion only checks the number of selected samples, which is not sufficient to verify that the diversity selection logic is working correctly. The test should be strengthened to check which samples are selected. You could add a unique ID to each sample in ds_list, and then after processing, verify that the IDs of the selected samples match a list of expected IDs for a given seed and input.