Updated tsne.py

Author: Nikita-Kedari

machine_learning/tsne.py

@@ -5,7 +5,7 @@
 t-SNE is a nonlinear dimensionality reduction algorithm for visualizing
 high-dimensional data in a low-dimensional space (2D or 3D).
 
-It computes pairwise similarities in both spaces and minimizes the 
+It computes pairwise similarities in both spaces and minimizes the
 Kullback-Leibler divergence using gradient descent.
 
 References:
@@ -149,35 +149,33 @@ def apply_tsne(
     return y
 
 
-def main() -> None:
+def main() -> ndarray:
     """
-    Run t-SNE on Iris dataset and display the first 5 embeddings.
+    Run t-SNE on Iris dataset and return the embeddings.
+
+    Returns:
+        ndarray: t-SNE embedding of the Iris dataset
 
     Example:
-    >>> main()  # doctest: +ELLIPSIS
-    t-SNE embedding (first 5 points):
-    [[-...
+    >>> result = main()
+    >>> result.shape
+    (150, 2)
+    >>> isinstance(result, np.ndarray)
+    True
     """
     data_x, _ = collect_dataset()
     y_emb = apply_tsne(data_x, n_components=2, n_iter=300)
 
     if not isinstance(y_emb, np.ndarray):
         raise TypeError("t-SNE embedding must be an ndarray")
 
-    print("t-SNE embedding (first 5 points):")
-    print(y_emb[:5])
-
-    # Optional visualization (commented, Ruff/mypy compliant)
-    # import matplotlib.pyplot as plt
-    # plt.scatter(
-    #     y_emb[:, 0],
-    #     y_emb[:, 1],
-    #     c=_labels,
-    #     cmap="viridis"
-    # )
-    # plt.show()
+    return y_emb
 
 
 if __name__ == "__main__":
     doctest.testmod()
-    main()
+
+    # Demonstration of the algorithm
+    result = main()
+    print("t-SNE embedding (first 5 points):")
+    print(result[:5])