Update README.md

README.md

@@ -110,35 +110,86 @@ This [Model2Vec](https://github.com/MinishLab/model2vec) model is a distilled ve
 
 ## Installation
 
-Install model2vec using pip:
+Install using pip:
 ```
-pip install model2vec
+pip install model2vec reach tqdm numpy
 ```
 
 ## Usage
-Load this model using the `from_pretrained` method:
-```python
-from model2vec import StaticModel
-
-# Load a pretrained Model2Vec model
-model = StaticModel.from_pretrained("jester6136/multilingual-e5-large-m2v")
-
-# Compute text embeddings
-embeddings = model.encode(["Example sentence"])
 ```
-
-Alternatively, you can distill your own model using the `distill` method:
-```python
-from model2vec.distill import distill
-
-# Choose a Sentence Transformer model
-model_name = "BAAI/bge-base-en-v1.5"
-
-# Distill the model
-m2v_model = distill(model_name=model_name, pca_dims=256)
-
-# Save the model
-m2v_model.save_pretrained("m2v_model")
+import numpy as np
+from model2vec import StaticModel
+from reach import Reach
+from tqdm import tqdm
+import time
+
+class TextDeduplicator:
+    def __init__(self, model_path: str):
+        # Load the pre-trained model
+        self.model = StaticModel.from_pretrained(model_path)
+
+    def encode_texts(self, texts: list[str]) -> np.ndarray:
+        # Prepare the texts and encode them into embeddings
+        texts = [f"query: {text}" for text in texts]
+        embedding_matrix = self.model.encode(texts, show_progressbar=True)
+        return embedding_matrix
+
+    def deduplicate(self, embedding_matrix: np.ndarray, threshold: float, batch_size: int = 1024
+    ) -> tuple[np.ndarray, dict[int, list[int]]]:
+        # Deduplicate the texts based on their embeddings
+        reach = Reach(vectors=embedding_matrix, items=[str(i) for i in range(len(embedding_matrix))])
+        results = reach.nearest_neighbor_threshold(
+            embedding_matrix, threshold=threshold, batch_size=batch_size, show_progressbar=True
+        )
+
+        deduplicated_indices = set(range(len(embedding_matrix)))
+        duplicate_groups = {}
+
+        for i, similar_items in enumerate(tqdm(results)):
+            if i not in deduplicated_indices:
+                continue
+
+            similar_indices = [int(item[0]) for item in similar_items if int(item[0]) != i]
+            for sim_idx in similar_indices:
+                if sim_idx in deduplicated_indices:
+                    deduplicated_indices.remove(sim_idx)
+                    if i not in duplicate_groups:
+                        duplicate_groups[i] = []
+                    duplicate_groups[i].append(sim_idx)
+
+        return np.array(list(deduplicated_indices)), duplicate_groups
+
+    def deduplicate_texts(self, texts: list[str], threshold: float) -> tuple[np.ndarray, dict[int, list[int]]]:
+        # End-to-end deduplication process
+        embedding_matrix = self.encode_texts(texts)
+        return self.deduplicate(embedding_matrix, threshold)
+
+
+if __name__ == "__main__":
+    # Example usage
+    texts = [
+        "Anh yêu em.",
+        "Mọi thứ ở công ty mới đều lạ lẫm, nhưng tôi cảm thấy rất sẵn sàng để bắt đầu hành trình mới.",
+        "Trận đấu bóng đá tối qua rất căng thẳng, hai đội liên tục tấn công và phòng thủ.",
+        "Một quan chức Fed muốn giảm bớt tốc độ hạ lãi suất",
+        "Ngày đầu tiên tại công ty mới đầy ấn tượng, tôi hy vọng sẽ nhanh chóng hòa nhập với môi trường làm việc.",
+        "Mùa hè này, cả gia đình sẽ có một chuyến đi đến Đà Nẵng, nơi mà chúng tôi đã mong chờ từ rất lâu.",
+        "Gia đình tôi đã lên kế hoạch cho kỳ nghỉ tại Đà Nẵng vào mùa hè này, một chuyến đi mà mọi người đều háo hức.",
+        "Fed có bước tiến mới để hạ lãi suất",
+        "Chúng tôi đã dự định từ lâu sẽ đi Đà Nẵng vào mùa hè này, và cả nhà đều rất trông đợi chuyến du lịch.",
+        "Ngày đầu đi làm thật là thú vị, tuy có chút hồi hộp nhưng tôi mong chờ những điều mới mẻ.",
+        "Mùa hè năm nay, gia đình tôi sẽ du lịch Đà Nẵng, chuyến đi mà ai cũng mong đợi từ trước."
+    ]
+
+    deduplicator = TextDeduplicator("jester6136/multilingual-e5-large-m2v")
+
+    start_time = time.time()
+    deduplicated_indices, duplicate_groups = deduplicator.deduplicate_texts(texts, threshold=0.85)
+    end_time = time.time()
+
+    print(f"Deduplication completed in {end_time - start_time:.2f} seconds")
+    print(f"Deduped output: {deduplicated_indices}")
+    print(f"Group dup: {duplicate_groups}")
 ```
 
 ## How it works