Text Segmentation: A Closer Look at Clustering and Beyond

import typing as t
from dataclasses import dataclass
import re
import warnings
warnings.filterwarnings("ignore")

from joblib import Memory

import pandas as pd
import numpy as np

import matplotlib.pyplot as plt
import plotly.express as px
import plotly.graph_objs as go
from plotly.subplots import make_subplots

from tqdm.notebook import tqdm

from stop_words import get_stop_words
from sentence_transformers import SentenceTransformer

memory = Memory('../.cache')

csv_path = '../data/400-415-episodes-ground-truth.csv'
df = pd.read_csv(csv_path)
df = df[df['episode'] == 412]
df.head()

@dataclass
class Models:
    paraphrase = 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2'
    labse = 'sentence-transformers/labse'

@memory.cache
def get_embeddings(sentences: list[str], model_name: str) -> list[float]:
    model = SentenceTransformer(model_name)
    return model.encode(sentences)

ru_paraphrase_embeddings = get_embeddings(df['ru_sentence'].values, Models.paraphrase)
en_paraphrase_embeddings = get_embeddings(df['en_sentence'].values, Models.paraphrase)

ru_labse_embeddings = get_embeddings(df['ru_sentence'].values, Models.labse)
en_labse_embeddings = get_embeddings(df['en_sentence'].values, Models.labse)

print(f'{ru_paraphrase_embeddings.shape=}')
print(f'{en_paraphrase_embeddings.shape=}')

print(f'{ru_labse_embeddings.shape=}')
print(f'{en_labse_embeddings.shape=}')

embeddings = {
    'russian': {
        'paraphrase': ru_paraphrase_embeddings,
        'labse': ru_labse_embeddings,
    },
    'english': {
        'paraphrase': en_paraphrase_embeddings,
        'labse': en_labse_embeddings,
    }
}

df['ground_truth_next'] = df['ground_truth'].shift(1).fillna(df['ground_truth'].max()).astype(int)

def is_first_sentence_in_topic(s: pd.Series) -> int:
    return 0 if s['ground_truth'] == s['ground_truth_next'] else 1

df['is_first_sentence_in_topic'] = df.apply(is_first_sentence_in_topic, axis='columns')

df[df['is_first_sentence_in_topic'] == 1][['ru_sentence', 'en_sentence', 'ground_truth', 'is_first_sentence_in_topic']]

color_columns = ['ground_truth', 'is_first_sentence_in_topic']

from sklearn.decomposition import PCA


def pca_scatter(embeddings: list[float], title: str, color_column: str) -> None:
    pca = PCA(n_components=2, random_state=42)
    pca_embeddings = pca.fit_transform(embeddings)

    pca_df = pd.DataFrame(pca_embeddings, columns=('x', 'y'))
    pca_df = df.join(pca_df.set_index(df.index))
    
    pca_df[color_column] = pca_df[color_column].astype(str)
    fig = px.scatter(
        title=f'PCA: {title}',
        data_frame=pca_df,
        x='x',
        y='y',
        color=color_column,
        hover_data=['ru_sentence', 'en_sentence'],
        width=1000
    )
    
    fig.update_layout(legend=dict(
        orientation='h',
        yanchor="bottom",
        y=1.02,
        xanchor="right",
        x=1
    ))
    
    return fig

for color_column in color_columns:
    for lang, embedding_pairs in embeddings.items():
        for model_name, embedding in embedding_pairs.items():
            pca_scatter(
                embeddings=embedding,
                title=f'lang={lang}, model_name={model_name}, color={color_column}',
                color_column=color_column).show(renderer='png')

from sklearn.manifold import TSNE
from scipy.spatial.distance import cosine


@memory.cache
def tsne_transform(embeddings: list[list[float]], metric: t.Union[str, t.Callable], early_exaggeration: int):
    tsne = TSNE(n_components=2, metric=metric, random_state=42, early_exaggeration=early_exaggeration, n_jobs=-1)
    tsne_embeddings = tsne.fit_transform(embeddings)
    return tsne_embeddings

def tsne_scatter(embeddings: list[list[float]], metric: t.Union[str, t.Callable], title: str, color_column: str) -> None:
    tsne_embeddings = tsne_transform(embeddings=embeddings, metric=metric, early_exaggeration=20)

    tsne_df = pd.DataFrame(tsne_embeddings, columns=('x', 'y'))
    tsne_df = df.join(tsne_df.set_index(df.index))    
    tsne_df[color_column] = tsne_df[color_column].astype(str)
    
    fig = px.scatter(
        title=f't-SNE: {title}',
        data_frame=tsne_df,
        x='x',
        y='y',
        color=color_column,
        hover_data=['ru_sentence', 'en_sentence'],
        width=1000)
    
    fig.update_layout(legend=dict(
        orientation='h',
        yanchor="bottom",
        y=1.02,
        xanchor="right",
        x=1
    ))
    
    return fig

metrics = {
    'cosine': cosine,
    'euclidean': 'euclidean',
}

for metric_name, metric in metrics.items():
    for color_column in color_columns:
        for lang, embedding_pairs in embeddings.items():
            for model_name, embedding in embedding_pairs.items():
                tsne_scatter(
                    embeddings=embedding,
                    metric=metric_name,
                    title=f'lang={lang}, metric={metric_name}, model_name={model_name}, color={color_column}',
                    color_column=color_column).show(renderer='png')

import umap


@memory.cache
def umap_transofrm(embeddings: list[list[float]], metric: t.Union[str, t.Callable]):
    umap_reducer = umap.UMAP(n_neighbors=15, min_dist=0.1, n_components=2, metric=metric)
    umap_embeddings = umap_reducer.fit_transform(embeddings)
    return umap_embeddings


def umap_scatter(embeddings: list[float], metric: t.Union[str, t.Callable], color_column: str, title: str):
    umap_embeddings = umap_transofrm(embeddings, metric)

    umap_df = pd.DataFrame(umap_embeddings, columns=('x', 'y'))
    umap_df = df.join(umap_df.set_index(df.index))    
    umap_df[color_column] = umap_df[color_column].astype(str)
        
    fig = px.scatter(
        title=f'UMAP: {title}',
        data_frame=umap_df,
        x='x',
        y='y',
        color=color_column,
        hover_data=['ru_sentence', 'en_sentence'],
        width=1000
    )

    fig.update_layout(legend=dict(
        orientation='h',
        yanchor="bottom",
        y=1.02,
        xanchor="right",
        x=1
    ))
    
    return fig

for metric_name, metric in metrics.items():
    for color_column in color_columns:
        for lang, embedding_pairs in embeddings.items():
            for model_name, embedding in embedding_pairs.items():
                umap_scatter(
                    embeddings=embedding,
                    metric=metric_name,
                    title=f'lang={lang}, metric={metric_name}, model_name={model_name}, color={color_column}',
                    color_column=color_column).show(renderer='png')

import plotly.graph_objs as go
from plotly.subplots import make_subplots


# Create data for the scatter plots
x1 = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
y1 = [10, 11, 12, 13, 14, 4, 7, 1, 5, 11]
z1 = [1, 2, 3, 4, 0, 5, 6, 7, 8, 9]

custom_colors = px.colors.qualitative.D3
color_scale = [custom_colors[val] for val in z1]

x2 = [2, 3, 4, 5, 6]
y2 = [8, 9, 10, 11, 12]

# Create scatter plot traces
trace1 = go.Scatter(
    x=x1,
    y=y1,
    name='Plot 1',
    mode='markers',
    marker=dict(
        color=color_scale,  # Color scale based on the Z axis
        colorscale='Viridis',  # Choose the desired colorscale (e.g., 'Viridis', 'Blues', 'reds', etc.)
    ))
trace2 = go.Scatter(x=x2, y=y2, mode='markers', name='Plot 2')

# Create a subplot
fig = make_subplots(rows=1, cols=2)

# Add the traces to the subplot
fig.add_trace(trace1, row=1, col=1)
fig.add_trace(trace2, row=1, col=2)

# Display the plot
fig.show()