Visualizing Sentence Embeddings with t-SNE in Java

Introduction

Welcome to the third lesson in our course on Text Representation Techniques for RAG systems! In our previous lesson, we explored how to generate sentence embeddings and saw how these richer representations capture semantic meaning better than the classic Bag-of-Words. Now, we will build on that knowledge to visualize these embeddings in a two-dimensional space using t-SNE (t-distributed Stochastic Neighbor Embedding). By the end of this lesson, you'll have an interactive way to see how thematically similar sentences group closer together, reinforcing the idea that embeddings preserve meaningful relationships between sentences.

Understanding t-SNE

t-SNE helps us visualize high-dimensional embeddings by compressing them into a given lower-dimensional space (usually 2D or 3D, for visualization) while preserving relative similarities : Similarity First : t-SNE prioritizes keeping similar sentences close. It calculates pairwise similarities in the original space (using a probability distribution) so nearby embeddings get higher similarity scores than distant ones. Local Structure : It preserves neighborhoods of related points rather than exact distances. This means clusters you see reflect genuine thematic groupings (e.g., NLP vs. Food), but axis values themselves have no intrinsic meaning. Perplexity Matters : This parameter (~5-50) controls neighborhood size. Lower values emphasize tight clusters (good for spotting subtopics), while higher values show broader trends (useful for separating major categories). Tradeoffs : While powerful for visualization, t-SNE is computationally expensive for large datasets (as it compares all sentence pairs). For RAG systems, this makes it better suited for exploratory analysis of smaller samples than production-scale data. You may be asking yourself, why does this matter for RAG? Seeing embeddings cluster by topic validates they're capturing semantic relationships — a prerequisite for effective retrieval. If NLP sentences scattered randomly, we'd question the embedding quality before even building the RAG pipeline, prompting us to reevaluate the choice of the embedding model.

Building Our Data

To demonstrate how t-SNE reveals natural groupings, we'll gather sentences on four different topics: NLP, ML, Food, and Weather. Then, we assign each sentence a category so we can later color-code and shape-code the points in our 2D visualization. Javaimport java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; public class DataBuilder { public static Map<String, Object> getSentencesAndCategories() { List<String> sentences = new ArrayList<>(); List<String> categories = new ArrayList<>(); // Topic: NLP sentences.add("RAG stands for Retrieval-Augmented Generation."); sentences.add("Retrieval is a crucial aspect of modern NLP systems."); sentences.add("Generating text with correct facts is challenging."); sentences.add("Large language models can generate coherent text."); sentences.add("GPT models have billions of parameters."); sentences.add("Natural Language Processing enables computers to understand human language."); sentences.add("Word embeddings capture semantic relationships between words."); sentences.add("Transformer architectures revolutionized NLP research."); for (int i = 0; i < 8; i++) categories.add("NLP"); // Topic: Machine Learning sentences.add("Machine learning benefits from large datasets."); sentences.add("Supervised learning requires labeled data."); sentences.add("Reinforcement learning is inspired by behavioral psychology."); sentences.add("Neural networks can learn complex functions."); sentences.add("Overfitting is a common problem in ML."); sentences.add("Unsupervised learning uncovers hidden patterns in data."); sentences.add("Feature engineering is critical for model performance."); sentences.add("Cross-validation helps in assessing model generalization."); for (int i = 0; i < 8; i++) categories.add("ML"); // Topic: Food sentences.add("Bananas are commonly used in smoothies."); sentences.add("Oranges are rich in vitamin C."); sentences.add("Pizza is a popular Italian dish."); sentences.add("Cooking pasta requires boiling water."); sentences.add("Chocolate can be sweet or bitter."); sentences.add("Fresh salads are a healthy and refreshing meal."); sentences.add("Sushi combines rice, fish, and seaweed in a delicate balance."); sentences.add("Spices can transform simple ingredients into gourmet dishes."); for (int i = 0; i < 8; i++) categories.add("Food"); // Topic: Weather sentences.add("It often rains in the Amazon rainforest."); sentences.add("Summers can be very hot in the desert."); sentences.add("Hurricanes form over warm ocean waters."); sentences.add("Snowstorms can disrupt transportation."); sentences.add("A sunny day can lift people's mood."); sentences.add("Foggy mornings are common in coastal regions."); sentences.add("Winter brings frosty nights and chilly winds."); sentences.add("Thunderstorms can produce lightning and heavy rain."); for (int i = 0; i < 8; i++) categories.add("Weather"); Map<String, Object> data = new HashMap<>(); data.put("sentences", sentences); data.put("categories", categories); return data; } public static Map<String, String> getColorAndShapeMaps() { Map<String, String> colorMap = new HashMap<>(); colorMap.put("NLP", "red"); colorMap.put("ML", "blue"); colorMap.put("Food", "green"); colorMap.put("Weather", "purple"); Map<String, String> shapeMap = new HashMap<>(); shapeMap.put("NLP", "o"); shapeMap.put("ML", "s"); shapeMap.put("Food", "^"); shapeMap.put("Weather", "X"); Map<String, String> maps = new HashMap<>(); maps.putAll(colorMap); maps.putAll(shapeMap); return maps; } }import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; public class DataBuilder { public static Map<String, Object> getSentencesAndCategories() { List<String> sentences = new ArrayList<>(); List<String> categories = new ArrayList<>(); // Topic: NLP sentences.add("RAG stands for Retrieval-Augmented Generation."); sentences.add("Retrieval is a crucial aspect of modern NLP systems."); sentences.add("Generating text with correct facts is challenging."); sentences.add("Large language models can generate coherent text."); sentences.add("GPT models have billions of parameters."); sentences.add("Natural Language Processing enables computers to understand human language."); sentences.add("Word embeddings capture semantic relationships between words."); sentences.add("Transformer architectures revolutionized NLP research."); for (int i = 0; i < 8; i++) categories.add("NLP"); // Topic: Machine Learning sentences.add("Machine learning benefits from large datasets."); sentences.add("Supervised learning requires labeled data."); sentences.add("Reinforcement learning is inspired by behavioral psychology."); sentences.add("Neural networks can learn complex functions."); sentences.add("Overfitting is a common problem in ML."); sentences.add("Unsupervised learning uncovers hidden patterns in data."); sentences.add("Feature engineering is critical for model performance."); sentences.add("Cross-validation helps in assessing model generalization."); for (int i = 0; i < 8; i++) categories.add("ML"); // Topic: Food sentences.add("Bananas are commonly used in smoothies."); sentences.add("Oranges are rich in vitamin C."); sentences.add("Pizza is a popular Italian dish."); sentences.add("Cooking pasta requires boiling water."); sentences.add("Chocolate can be sweet or bitter."); sentences.add("Fresh salads are a healthy and refreshing meal."); sentences.add("Sushi combines rice, fish, and seaweed in a delicate balance."); sentences.add("Spices can transform simple ingredients into gourmet dishes."); for (int i = 0; i < 8; i++) categories.add("Food"); // Topic: Weather sentences.add("It often rains in the Amazon rainforest."); sentences.add("Summers can be very hot in the desert."); sentences.add("Hurricanes form over warm ocean waters."); sentences.add("Snowstorms can disrupt transportation."); sentences.add("A sunny day can lift people's mood."); sentences.add("Foggy mornings are common in coastal regions."); sentences.add("Winter brings frosty nights and chilly winds."); sentences.add("Thunderstorms can produce lightning and heavy rain."); for (int i = 0; i < 8; i++) categories.add("Weather"); Map<String, Object> data = new HashMap<>(); data.put("sentences", sentences); data.put("categories", categories); return data; } public static Map<String, String> getColorAndShapeMaps() { Map<String, String> colorMap = new HashMap<>(); colorMap.put("NLP", "red"); colorMap.put("ML", "blue"); colorMap.put("Food", "green"); colorMap.put("Weather", "purple"); Map<String, String> shapeMap = new HashMap<>(); shapeMap.put("NLP", "o"); shapeMap.put("ML", "s"); shapeMap.put("Food", "^"); shapeMap.put("Weather", "X"); Map<String, String> maps = new HashMap<>(); maps.putAll(colorMap); maps.putAll(shapeMap); return maps; } } Here's what's going on: The getSentencesAndCategories method returns a map containing two lists: one with sentences, another labeling each sentence's category. The getColorAndShapeMaps method creates two maps that tell the plotting function which colors and marker shapes to use per category (e.g., “red circles” for NLP).

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Java

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class DataBuilder {
    public static Map<String, Object> getSentencesAndCategories() {
        List<String> sentences = new ArrayList<>();
        List<String> categories = new ArrayList<>();

        // Topic: NLP
        sentences.add("RAG stands for Retrieval-Augmented Generation.");
        sentences.add("Retrieval is a crucial aspect of modern NLP systems.");
        sentences.add("Generating text with correct facts is challenging.");
        sentences.add("Large language models can generate coherent text.");
        sentences.add("GPT models have billions of parameters.");
        sentences.add("Natural Language Processing enables computers to understand human language.");
        sentences.add("Word embeddings capture semantic relationships between words.");
        sentences.add("Transformer architectures revolutionized NLP research.");
        for (int i = 0; i < 8; i++) categories.add("NLP");

        // Topic: Machine Learning
        sentences.add("Machine learning benefits from large datasets.");
        sentences.add("Supervised learning requires labeled data.");
        sentences.add("Reinforcement learning is inspired by behavioral psychology.");
        sentences.add("Neural networks can learn complex functions.");
        sentences.add("Overfitting is a common problem in ML.");
        sentences.add("Unsupervised learning uncovers hidden patterns in data.");
        sentences.add("Feature engineering is critical for model performance.");
        sentences.add("Cross-validation helps in assessing model generalization.");
        for (int i = 0; i < 8; i++) categories.add("ML");

        // Topic: Food
        sentences.add("Bananas are commonly used in smoothies.");
        sentences.add("Oranges are rich in vitamin C.");
        sentences.add("Pizza is a popular Italian dish.");
        sentences.add("Cooking pasta requires boiling water.");
        sentences.add("Chocolate can be sweet or bitter.");
        sentences.add("Fresh salads are a healthy and refreshing meal.");
        sentences.add("Sushi combines rice, fish, and seaweed in a delicate balance.");
        sentences.add("Spices can transform simple ingredients into gourmet dishes.");
        for (int i = 0; i < 8; i++) categories.add("Food");

        // Topic: Weather
        sentences.add("It often rains in the Amazon rainforest.");
        sentences.add("Summers can be very hot in the desert.");
        sentences.add("Hurricanes form over warm ocean waters.");
        sentences.add("Snowstorms can disrupt transportation.");
        sentences.add("A sunny day can lift people's mood.");
        sentences.add("Foggy mornings are common in coastal regions.");
        sentences.add("Winter brings frosty nights and chilly winds.");
        sentences.add("Thunderstorms can produce lightning and heavy rain.");
        for (int i = 0; i < 8; i++) categories.add("Weather");

        Map<String, Object> data = new HashMap<>();
        data.put("sentences", sentences);
        data.put("categories", categories);
        return data;
    }

    public static Map<String, String> getColorAndShapeMaps() {
        Map<String, String> colorMap = new HashMap<>();
        colorMap.put("NLP", "red");
        colorMap.put("ML", "blue");
        colorMap.put("Food", "green");
        colorMap.put("Weather", "purple");

        Map<String, String> shapeMap = new HashMap<>();
        shapeMap.put("NLP", "o");
        shapeMap.put("ML", "s");
        shapeMap.put("Food", "^");
        shapeMap.put("Weather", "X");

        Map<String, String> maps = new HashMap<>();
        maps.putAll(colorMap);
        maps.putAll(shapeMap);
        return maps;
    }
}

Java

import com.robrua.nlp.easybert.Bert;
import org.nd4j.linalg.api.ndarray.INDArray;
import smile.manifold.TSNE;

import java.util.List;
import java.util.stream.Collectors;

public class EmbeddingReducer {
    public static double[][] computeTsneEmbeddings(List<String> sentences, int perplexity, int iterations) throws Exception {
        // 1. Load the easyBert model.
        Bert bert = Bert.load("path/to/bert/model");

        // 2. Convert each sentence into a high-dimensional embedding.
        List<INDArray> embeddings = sentences.stream()
                .map(sentence -> bert.embedSequence(sentence))
                .collect(Collectors.toList());

        // 3. Convert the list of INDArrays to a 2D array for t-SNE.
        double[][] embeddingArray = new double[embeddings.size()][];
        for (int i = 0; i < embeddings.size(); i++) {
            embeddingArray[i] = embeddings.get(i).toDoubleVector();
        }

        // 4. Configure t-SNE with chosen parameters and reduce embeddings to 2D.
        TSNE tsne = new TSNE.Builder()
                .setPerplexity(perplexity)
                .setMaxIter(iterations)
                .build();
        return tsne.fit(embeddingArray);
    }
}

import com.robrua.nlp.easybert.Bert;
import org.nd4j.linalg.api.ndarray.INDArray;
import smile.manifold.TSNE;

import java.util.List;
import java.util.stream.Collectors;

public class EmbeddingReducer {
    public static double[][] computeTsneEmbeddings(List<String> sentences, int perplexity, int iterations) throws Exception {
        // 1. Load the easyBert model.
        Bert bert = Bert.load("path/to/bert/model");

        // 2. Convert each sentence into a high-dimensional embedding.
        List<INDArray> embeddings = sentences.stream()
                .map(sentence -> bert.embedSequence(sentence))
                .collect(Collectors.toList());

        // 3. Convert the list of INDArrays to a 2D array for t-SNE.
        double[][] embeddingArray = new double[embeddings.size()][];
        for (int i = 0; i < embeddings.size(); i++) {
            embeddingArray[i] = embeddings.get(i).toDoubleVector();
        }

        // 4. Configure t-SNE with chosen parameters and reduce embeddings to 2D.
        TSNE tsne = new TSNE.Builder()
                .setPerplexity(perplexity)
                .setMaxIter(iterations)
                .build();
        return tsne.fit(embeddingArray);
    }
}

Java

import matplotlib4j.plot.Figure;
import matplotlib4j.plot.Plot;
import java.util.List;
import java.util.Map;

public class EmbeddingPlotter {
    public static void plotEmbeddings(double[][] reducedEmbeddings, List<String> sentences,
                                      List<String> categories, Map<String, String> colorMap,
                                      Map<String, String> shapeMap) {
        Figure fig = new Figure();
        Plot plt = fig.plot();

        // 1. Plot each category with its corresponding color and shape.
        for (String category : colorMap.keySet()) {
            double[] x = new double[sentences.size()];
            double[] y = new double[sentences.size()];
            int index = 0;
            for (int i = 0; i < sentences.size(); i++) {
                if (categories.get(i).equals(category)) {
                    x[index] = reducedEmbeddings[i][0];
                    y[index] = reducedEmbeddings[i][1];
                    index++;
                }
            }
            plt.scatter(x, y, colorMap.get(category), shapeMap.get(category), category);
        }

        // 2. Set plot labels and title.
        plt.xlabel("t-SNE Dimension 1");
        plt.ylabel("t-SNE Dimension 2");
        plt.title("t-SNE Visualization of Sentence Embeddings");

        // 3. Show the plot.
        fig.show();
    }
}

import matplotlib4j.plot.Figure;
import matplotlib4j.plot.Plot;
import java.util.List;
import java.util.Map;

public class EmbeddingPlotter {
    public static void plotEmbeddings(double[][] reducedEmbeddings, List<String> sentences,
                                      List<String> categories, Map<String, String> colorMap,
                                      Map<String, String> shapeMap) {
        Figure fig = new Figure();
        Plot plt = fig.plot();

        // 1. Plot each category with its corresponding color and shape.
        for (String category : colorMap.keySet()) {
            double[] x = new double[sentences.size()];
            double[] y = new double[sentences.size()];
            int index = 0;
            for (int i = 0; i < sentences.size(); i++) {
                if (categories.get(i).equals(category)) {
                    x[index] = reducedEmbeddings[i][0];
                    y[index] = reducedEmbeddings[i][1];
                    index++;
                }
            }
            plt.scatter(x, y, colorMap.get(category), shapeMap.get(category), category);
        }

        // 2. Set plot labels and title.
        plt.xlabel("t-SNE Dimension 1");
        plt.ylabel("t-SNE Dimension 2");
        plt.title("t-SNE Visualization of Sentence Embeddings");

        // 3. Show the plot.
        fig.show();
    }
}

Java

import com.robrua.nlp.easybert.Bert;
import org.nd4j.linalg.api.ndarray.INDArray;
import smile.manifold.TSNE;

import java.util.List;
import java.util.stream.Collectors;

public class EmbeddingReducer {
    public static double[][] computeTsneEmbeddings(List<String> sentences, int perplexity, int iterations) throws Exception {
        // 1. Load the easyBert model.
        Bert bert = Bert.load("path/to/bert/model");

        // 2. Convert each sentence into a high-dimensional embedding.
        List<INDArray> embeddings = sentences.stream()
                .map(sentence -> bert.embedSequence(sentence))
                .collect(Collectors.toList());

        // 3. Convert the list of INDArrays to a 2D array for t-SNE.
        double[][] embeddingArray = new double[embeddings.size()][];
        for (int i = 0; i < embeddings.size(); i++) {
            embeddingArray[i] = embeddings.get(i).toDoubleVector();
        }

        // 4. Configure t-SNE with chosen parameters and reduce embeddings to 2D.
        TSNE tsne = new TSNE.Builder()
                .setPerplexity(perplexity)
                .setMaxIter(iterations)
                .build();
        return tsne.fit(embeddingArray);
    }
}

Java

import matplotlib4j.plot.Figure;
import matplotlib4j.plot.Plot;
import java.util.List;
import java.util.Map;

public class EmbeddingPlotter {
    public static void plotEmbeddings(double[][] reducedEmbeddings, List<String> sentences,
                                      List<String> categories, Map<String, String> colorMap,
                                      Map<String, String> shapeMap) {
        Figure fig = new Figure();
        Plot plt = fig.plot();

        // 1. Plot each category with its corresponding color and shape.
        for (String category : colorMap.keySet()) {
            double[] x = new double[sentences.size()];
            double[] y = new double[sentences.size()];
            int index = 0;
            for (int i = 0; i < sentences.size(); i++) {
                if (categories.get(i).equals(category)) {
                    x[index] = reducedEmbeddings[i][0];
                    y[index] = reducedEmbeddings[i][1];
                    index++;
                }
            }
            plt.scatter(x, y, colorMap.get(category), shapeMap.get(category), category);
        }

        // 2. Set plot labels and title.
        plt.xlabel("t-SNE Dimension 1");
        plt.ylabel("t-SNE Dimension 2");
        plt.title("t-SNE Visualization of Sentence Embeddings");

        // 3. Show the plot.
        fig.show();
    }
}

Introduction

Understanding t-SNE

Building Our Data

Join the 1M+ learners on CodeSignal

Generating and Reducing Embeddings

Plotting the Embeddings

Interpreting the t-SNE Plot

Conclusion and Next Steps