TextEmbeddingPlot

Create a plot of the dimension-reduced text embeddings of a collection of strings

DefinitionResourceCreateNotebook

Clear[TextEmbeddingPlot]

TextEmbeddingPlot[data_/;ListQ[data],dim_/;dim===2||dim===3,opts:OptionsPattern[{Frame->True,Axes->False,FrameTicks->None,Ticks->None,LabelingFunction->Tooltip,"EmbeddingFunction"->ResourceFunction["SentenceBERTEmbedding"],"ReductionMethod"->Automatic,ListPlot,ListPointPlot3D}]]:=Module[
{embeddings,reducedEmbeddings,embedder,reducer},
embedder=OptionValue["EmbeddingFunction"];
If[Depth[data]===3,embeddings=Map[embedder,data],embeddings=embedder[data]];
Switch[
OptionValue["ReductionMethod"],
Alternatives@@{Automatic,"AutoEncoder","Hadamard","Isomap","LatentSemanticAnalysis","Linear","LLE","MultidimensionalScaling","PrincipalComponentsAnalysis","TSNE","UMAP"},
reducer=DimensionReduction[Flatten[Normal@embeddings,If[Depth[data]===3,1,0]],dim,Method->OptionValue["ReductionMethod"]];,
_DimensionReducerFunction,
reducer=OptionValue["ReductionMethod"];,
_,
reducer=If[Depth[data]===3,Map[OptionValue["ReductionMethod"]],OptionValue["ReductionMethod"]];
];
reducedEmbeddings=Map[reducer,Normal@embeddings];
Switch[
dim,
2,
ListPlot[reducedEmbeddings->data,FilterRules[DeleteDuplicatesBy[Join[{opts},{LabelingFunction->OptionValue[LabelingFunction],Frame->OptionValue[Frame],Axes->OptionValue[Axes],FrameTicks->OptionValue[FrameTicks]}],Last],Options@ListPlot]],
3,
ListPointPlot3D[reducedEmbeddings->data,FilterRules[DeleteDuplicatesBy[Join[{opts},{LabelingFunction->OptionValue[LabelingFunction],Ticks->OptionValue[Ticks],Axes->OptionValue[Axes]}],Last],Options@ListPointPlot3D]],
_,
$Failed
]
]

Documentation

UsageResourceCreateNotebook

Details & OptionsResourceCreateNotebook

The argument dim can be either 2 or 3.

By default, embedding vectors are computed using ResourceFunction["SentenceBERTEmbedding"].

By default, embedding vectors are computed with "Real32" working precision.

A Graphics object is returned.

The option "EmbeddingFunction" can be an arbitrary function which takes a list of strings as input and outputs a list of corresponding vectors.

The option "ReductionMethod" can be set to any of the Method options accepted by DimensionReduction. This option can also be set to a specific DimensionReducerFunction with appropriate input and output dimensions. Alternatively, this option can be set to an arbitrary function which accepts vectors of the appropriate input and output dimension.

The option "EmbeddingFunction" can be an arbitrary function which takes a list of strings as input and outputs a list of corresponding vectors.

Accepts the same options as ListPlot and ListPointPlot3D.

The default LabelingFunction is set to Tooltip, with the original strings used as labels.

ExamplesResourceCreateNotebook

Basic Examples

Plot the embeddings of a list of strings:

TextEmbeddingPlot[Alphabet[],2]

Compare sentence embeddings from different texts:

TextEmbeddingPlot[TextSentences/@{ResourceData["Gettysburg Address"],ResourceData["Friends, Romans, Countrymen"]},3]

Scope

Prepare sentences from three different text sources:

sources={"Alice in Wonderland","Declaration of Independence","Gettysburg Address"};
data=Map[Take[TextSentences@ResourceData[#],UpTo[50]]&,sources];

Create a labeled plot:

TextEmbeddingPlot[data,3,PlotLegends->sources]

Present the same data with a different dimension reduction method:

TextEmbeddingPlot[data,3,PlotLegends->sources,"ReductionMethod"->"UMAP"]

Options

"EmbeddingFunction"

Define a custom embedding function using an external service, which also prints information on model and token use:

customEmbedder[stringlist_List]/;AllTrue[stringlist,StringQ]:=
Module[{response},
response=ServiceExecute["OpenAI","Embedding",{"Input"->stringlist}];
Print@KeyDrop[response,"Content"];
response["Content"]
]

Use the custom embedding function on the previously defined data:

TextEmbeddingPlot[data[[All,1;;10]],2,"EmbeddingFunction"->customEmbedder]

The "EmbeddingFunction" should accept a list of strings as input. Use a naive embedding function that computes letter counts:

TextEmbeddingPlot[data,2,"EmbeddingFunction"->Map[Lookup[LetterCounts[#],Alphabet[],0]&]]

A function specified for "EmbeddingFunction" should output a corresponding list of vectors or NumericArray objects:

randomEmbedder[list_]:=NumericArray[RandomReal[{-1,1},{Length@list,384}],"Real32"]

TextEmbeddingPlot[data,2,"EmbeddingFunction"->randomEmbedder]

"ReductionMethod"

Method options of DimensionReduction are supported:

GraphicsRow@Table[TextEmbeddingPlot[data,2,"ReductionMethod"->m],{m,{Automatic,"TSNE","UMAP"}}]

Use a specific dimension reduction function:

dr=DimensionReduction[RandomReal[{-100,100},{1000,384}],2]

TextEmbeddingPlot[data,2,"ReductionMethod"->dr]

Define an arbitrary function to perform the dimension reduction on embedding vectors:

TextEmbeddingPlot[data,2,"ReductionMethod"->(Take[#,-2]&)]

Labeling Function

Use default behavior of ListPlot:

TextEmbeddingPlot[data,2,LabelingFunction->Automatic]

Use default behavior of ListPointPlot3D:

TextEmbeddingPlot[data,3,LabelingFunction->Automatic]

Define a custom LabelingFunction:

TextEmbeddingPlot[data[[All,1;;10]],2,LabelingFunction->(StringTake[#3[[3,1]],UpTo[20]]<>" ..."&)]

Applications

Properties and Relations

Possible Issues

A "ReductionMethod" that is not one of the supported options for DimensionReduction must accept vectors of the appropriate input and output dimensions:

wrongOutputDim=DimensionReduction[RandomReal[{-100,100},{1000,384}],3]

TextEmbeddingPlot[data,2,"ReductionMethod"->wrongOutputDim]

wrongInputDim=DimensionReduction[RandomReal[{-100,100},{1000,15}],2]

Enclose@ConfirmQuiet@TextEmbeddingPlot[data,2,"ReductionMethod"->wrongInputDim]

Neat Examples

Source & Additional Information

Contributed ByResourceCreateNotebook

John McNally

KeywordsResourceCreateNotebook

SentenceBERT

text embeddings

embedding visualization

natural language processing

nlp

CategoriesResourceCreateNotebook

Related SymbolsResourceCreateNotebook

Related Resource ObjectsResourceCreateNotebook

SentenceBERTEmbedding

Source/Reference CitationResourceCreateNotebook

Source, reference or citation information

LinksResourceCreateNotebook

Link to other related material

TestsResourceCreateNotebook

MyFunction[x,y]

CompatibilityResourceCreateNotebook

Wolfram Language VersionResourceCreateNotebook

14.0+

Operating SystemResourceCreateNotebook

EnvironmentsResourceCreateNotebook

Cloud SupportResourceCreateNotebook

Required FeaturesResourceCreateNotebook

Author NotesResourceCreateNotebook

Additional information about limitations, issues, etc.

Submission NotesResourceCreateNotebook

Additional information for the reviewer.