Embeddings Use Cases

# Ensure you have your API key set in your environment per the README: https://github.com/openai/openai-python#usage
import openai
openai.api_key = '[openai_key]'

import pandas as pd
import numpy as np
import tiktoken
import sys
from typing import List, Optional
from sklearn.manifold import TSNE
from ast import literal_eval

# Now try to import your module again
def get_embedding(text: str, engine="text-similarity-davinci-001", **kwargs) -> List[float]:

    # replace newlines, which can negatively affect performance.
    text = text.replace("\n", " ")

    return openai.Embedding.create(input=[text], engine=engine, **kwargs)["data"][0]["embedding"]

# Embedding model parameters
embedding_model = "text-embedding-ada-002"
embedding_encoding = "cl100k_base"  # this the encoding for text-embedding-ada-002
max_tokens = 8000  # the maximum for text-embedding-ada-002 is 8191

from datasets import load_dataset

dataset = load_dataset("rotten_tomatoes")
df = dataset['train'].to_pandas()
print("{x} rows in df.".format(x=len(df)))

top_n = 5000

encoding = tiktoken.get_encoding(embedding_encoding)

# omit reviews that are too long to embed
df["n_tokens"] = df['text'].apply(lambda x: len(encoding.encode(x)))
df = df[df.n_tokens <= max_tokens].tail(top_n)

# get embeddings and save them
df["embedding"] = df['text'].apply(lambda x: get_embedding(x, engine=embedding_model))
df.to_csv("./rotten_tomatoes_with_embeddings_{x}.csv".format(x=top_n))

import numpy as np
from sklearn.cluster import KMeans

# Convert to a list of lists of floats
matrix = np.vstack(df.embedding.apply(literal_eval).to_list())

n_clusters = 5

kmeans = KMeans(n_clusters=n_clusters, init='k-means++', random_state=42)
kmeans.fit(matrix)
df['cluster'] = kmeans.labels_

df.groupby("cluster").label.mean().sort_values()

from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
import plotly.express as px

tsne = TSNE(n_components=2, perplexity=15, random_state=42, init="random", learning_rate=200)
vis_dims2 = tsne.fit_transform(matrix)

# Define the color map for the clusters
color_map = {
    0: 'firebrick',         # Cluster 0
    1: 'orangered',  # Cluster 1
    2: 'gold',      # Cluster 2
    3: 'limegreen',       # Cluster 3
    4: 'maroon'       # Cluster 4
}

# Map the cluster labels to colors
df['color'] = df['cluster'].map(color_map)

# Convert cluster to a categorical type with the specified order
df['cluster'] = pd.Categorical(df['cluster'], categories=[0, 1, 2, 3, 4], ordered=True)

# Create a DataFrame from the t-SNE results
df_tsne = pd.DataFrame({'Dimension 1': vis_dims2[:, 0], 
                        'Dimension 2': vis_dims2[:, 1], 
                        'Cluster': df['cluster'],
                        'Color': df['color']})

# Create the 2D scatter plot using Plotly Express
fig = px.scatter(
    df_tsne, x='Dimension 1', y='Dimension 2',
    color='Cluster', labels={'Cluster': 'Cluster'},
    title='Clusters identified visualized in 2D using t-SNE',
    color_discrete_map=color_map # Apply the color map
)

# Define the category orders for the legend to make it discrete
fig.update_traces(marker=dict(size=3), selector=dict(mode='markers'))  # Adjust marker size
fig.update_layout(
    legend=dict(
        traceorder='normal',
        title_text='Cluster',
        title_font=dict(size=14)
    )
)


# Show the plot
fig.show()

Cluster 0 Theme:  All of the reviews are negative and express dissatisfaction with the product or experience.
1, now as a former gong show addict , i'll admit it , my only complaint i
0, there's just no currency in deriding james bond for being a clichéd , 
0, ecks this one off your must-see list .
0, skip this turd and pick your nose instead because you're sure to get m
0, this movie is about the worst thing chan has done in the united states
----------------------------------------------------------------------------------------------------
Cluster 1 Theme:  The reviews are all negative and critical of the movie.
0, the type of dumbed-down exercise in stereotypes that gives the [teen c
0, just not campy enough
0, not so much farcical as sour .
0, a one-trick pony whose few t&a bits still can't save itself from being
0, cuba gooding jr . valiantly mugs his way through snow dogs , but even 
----------------------------------------------------------------------------------------------------
Cluster 2 Theme:  Disappointed with the quality of the movie.
0, its generic villains lack any intrigue ( other than their funny accent
0, one of the most highly-praised disappointments i've had the misfortune
0, . . . with the candy-like taste of it fading faster than 25-cent bubbl
0, for all its impressive craftsmanship , and despite an overbearing seri
0, the script by vincent r . nebrida . . . tries to cram too many ingredi
----------------------------------------------------------------------------------------------------
Cluster 3 Theme:  All of the reviews are positive and praise the movie for its unique qualities, such as its surreal sense of humor, technological finish, insightful writing, delicate performances, and character-driven storytelling.
1, what elevates the movie above the run-of-the-mill singles blender is i
0, at least it's a fairly impressive debut from the director , charles st
1, insightfully written , delicately performed
1, one of those exceedingly rare films in which the talk alone is enough 
1, a stylish but steady , and ultimately very satisfying , piece of chara
----------------------------------------------------------------------------------------------------
Cluster 4 Theme:  Disappointment with the quality of the product or experience.
0, qualities that were once amusing are becoming irritating .
0, everything's serious , poetic , earnest and -- sadly -- dull .
1, what's infuriating about full frontal is that it's too close to real l
0, befuddled in its characterizations as it begins to seem as long as the
0, human nature talks the talk , but it fails to walk the silly walk that
----------------------------------------------------------------------------------------------------

# Reading a review which belong to each group.
rev_per_cluster = n_clusters

for i in range(n_clusters):
    print(f"Cluster {i} Theme:", end=" ")

    reviews = "\n".join(
        df[df['cluster'] == i]
        .text.str.replace("Title: ", "")
        .str.replace("\n\nContent: ", ":  ")
        .sample(rev_per_cluster, random_state=42)
        .values
    )
    response = openai.Completion.create(
        engine="text-davinci-003",
        prompt=f'What do the following customer reviews have in common?\n\nCustomer reviews:\n"""\n{reviews}\n"""\n\nTheme:',
        temperature=0,
        max_tokens=64,
        top_p=1,
        frequency_penalty=0,
        presence_penalty=0,
    )
    print(response["choices"][0]["text"].replace("\n", ""))

    sample_cluster_rows = df[df.cluster == i].sample(rev_per_cluster, random_state=42)
    for j in range(rev_per_cluster):
        print(sample_cluster_rows.label.values[j], end=", ")
        print(sample_cluster_rows.text.str[:70].values[j])

    print("-" * 100)