Agrupación con embeddings

# @title Licensed under the Apache License, Version 2.0 (the "License");
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

%pip install -U -q google-genai

import re
import tqdm
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

from google import genai
from google.genai import types

# Used to securely store your API key
from google.colab import userdata

from sklearn.datasets import fetch_20newsgroups
from sklearn.manifold import TSNE
from sklearn.cluster import KMeans
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

# Or use `os.getenv('GEMINI_API_KEY')` to fetch an environment variable.
GEMINI_API_KEY=userdata.get('GEMINI_API_KEY')

client = genai.Client(api_key=GEMINI_API_KEY)

for m in client.models.list():
  if 'embedContent' in m.supported_actions:
    print(m.name)

models/embedding-001
models/text-embedding-004
models/gemini-embedding-exp-03-07
models/gemini-embedding-exp
models/gemini-embedding-001

MODEL_ID = "gemini-embedding-001" # @param ["gemini-embedding-2-preview", "gemini-embedding-001"] {"allow-input":true, isTemplate: true}

newsgroups_train = fetch_20newsgroups(subset='train')

# View list of class names for dataset
newsgroups_train.target_names

['alt.atheism',
 'comp.graphics',
 'comp.os.ms-windows.misc',
 'comp.sys.ibm.pc.hardware',
 'comp.sys.mac.hardware',
 'comp.windows.x',
 'misc.forsale',
 'rec.autos',
 'rec.motorcycles',
 'rec.sport.baseball',
 'rec.sport.hockey',
 'sci.crypt',
 'sci.electronics',
 'sci.med',
 'sci.space',
 'soc.religion.christian',
 'talk.politics.guns',
 'talk.politics.mideast',
 'talk.politics.misc',
 'talk.religion.misc']

idx = newsgroups_train.data[0].index('Lines')
print(newsgroups_train.data[0][idx:])

Lines: 15

 I was wondering if anyone out there could enlighten me on this car I saw
the other day. It was a 2-door sports car, looked to be from the late 60s/
early 70s. It was called a Bricklin. The doors were really small. In addition,
the front bumper was separate from the rest of the body. This is 
all I know. If anyone can tellme a model name, engine specs, years
of production, where this car is made, history, or whatever info you
have on this funky looking car, please e-mail.

Thanks,
- IL
   ---- brought to you by your neighborhood Lerxst ----

# Apply functions to remove names, emails, and extraneous words from data points in newsgroups.data
newsgroups_train.data = [re.sub(r'[\w\.-]+@[\w\.-]+', '', d) for d in newsgroups_train.data] # Remove email
newsgroups_train.data = [re.sub(r"\([^()]*\)", "", d) for d in newsgroups_train.data] # Remove names
newsgroups_train.data = [d.replace("From: ", "") for d in newsgroups_train.data] # Remove "From: "
newsgroups_train.data = [d.replace("\nSubject: ", "") for d in newsgroups_train.data] # Remove "\nSubject: "

# Put training points into a dataframe
df_train = pd.DataFrame(newsgroups_train.data, columns=['Text'])
df_train['Label'] = newsgroups_train.target
# Match label to target name index
df_train['Class Name'] = df_train['Label'].map(newsgroups_train.target_names.__getitem__)
# Retain text samples that can be used in the gecko model.
df_train = df_train[df_train['Text'].str.len() < 10000]

df_train

Text  Label  \
0       WHAT car is this!?\nNntp-Posting-Host: rac3.w...      7   
1       SI Clock Poll - Final Call\nSummary: Final ca...      4   
2       PB questions...\nOrganization: Purdue Univers...      4   
3       Re: Weitek P9000 ?\nOrganization: Harris Comp...      1   
4       Re: Shuttle Launch Question\nOrganization: Sm...     14   
...                                                  ...    ...   
11309    Re: Migraines and scans\nDistribution: world...     13   
11310  Screen Death: Mac Plus/512\nLines: 22\nOrganiz...      4   
11311   Mounting CPU Cooler in vertical case\nOrganiz...      3   
11312   Re: Sphere from 4 points?\nOrganization: Cent...      1   
11313   stolen CBR900RR\nOrganization: California Ins...      8   

                     Class Name  
0                     rec.autos  
1         comp.sys.mac.hardware  
2         comp.sys.mac.hardware  
3                 comp.graphics  
4                     sci.space  
...                         ...  
11309                   sci.med  
11310     comp.sys.mac.hardware  
11311  comp.sys.ibm.pc.hardware  
11312             comp.graphics  
11313           rec.motorcycles  

[11141 rows x 3 columns]

# Take a sample of each label category from df_train
SAMPLE_SIZE = 150
df_train = (df_train.groupby('Label', as_index = False)
                    .apply(lambda x: x.sample(SAMPLE_SIZE))
                    .reset_index(drop=True))

# Choose categories about science
df_train = df_train[df_train['Class Name'].str.contains('sci')]

# Reset the index
df_train = df_train.reset_index()
df_train

/tmp/ipykernel_200977/406673449.py:4: FutureWarning: DataFrameGroupBy.apply operated on the grouping columns. This behavior is deprecated, and in a future version of pandas the grouping columns will be excluded from the operation. Either pass `include_groups=False` to exclude the groupings or explicitly select the grouping columns after groupby to silence this warning.
  .apply(lambda x: x.sample(SAMPLE_SIZE))

index                                               Text  Label  \
0     1650   (Stephan Neuhaus )Re: PGP 2.2: general commen...     11   
1     1651   Re: Off the shelf cheap DES keyseach machine ...     11   
2     1652   Re: text of White House announcement and Q&As...     11   
3     1653   Need help !!\nKeywords: Firewall gateway mode...     11   
4     1654   Re: text of White House announcement and Q&As...     11   
..     ...                                                ...    ...   
595   2245   Re: Space Station Redesign, JSC Alternative #...     14   
596   2246   Portuguese Launch Complex \nOrganization: NAS...     14   
597   2247  Subject: <None>\n\nOrganization: University of...     14   
598   2248   Re: Vandalizing the sky.\nNntp-Posting-Host: ...     14   
599   2249   Re: Crazy? or just Imaginitive?\nOrganization...     14   

    Class Name  
0    sci.crypt  
1    sci.crypt  
2    sci.crypt  
3    sci.crypt  
4    sci.crypt  
..         ...  
595  sci.space  
596  sci.space  
597  sci.space  
598  sci.space  
599  sci.space  

[600 rows x 4 columns]

df_train['Class Name'].value_counts()

Class Name
sci.crypt          150
sci.electronics    150
sci.med            150
sci.space          150
Name: count, dtype: int64

Generar los embeddings

En esta sección, verás cómo generar embeddings para los diferentes textos en el dataframe usando los embeddings de la API de Gemini.

El modelo de embedding de Gemini admite varios tipos de tareas, cada uno adaptado a un objetivo específico. Aquí tienes una descripción general de los tipos disponibles y sus aplicaciones:

from tqdm.auto import tqdm
tqdm.pandas()

from google.api_core import retry

def make_embed_text_fn(model):

  @retry.Retry(timeout=300.0)
  def embed_fn(text: str) -> list[float]:
    # Set the task_type to CLUSTERING.
    result = client.models.embed_content(model=model,
                                            contents=text,
                                            config=types.EmbedContentConfig(
                                                task_type="clustering"))
    return np.array(result.embeddings[0].values)

  return embed_fn

def create_embeddings(df):
  df['Embeddings'] = df['Text'].progress_apply(make_embed_text_fn(MODEL_ID))
  return df

df_train = create_embeddings(df_train)

/home/lucianomartins/Documents/projects/gemini-launches/gemini-embedding/ge-env/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm
100%|██████████| 600/600 [05:15<00:00,  1.90it/s]

Reducción de dimensionalidad

La dimensión del vector de embedding del documento es 3072. Para visualizar cómo se agrupan los documentos embebidos, deberás aplicar la reducción de dimensionalidad, ya que solo puedes visualizar los embeddings en un espacio 2D o 3D. Los documentos contextualmente similares deben estar más cerca en el espacio que los documentos que no son tan similares.

len(df_train['Embeddings'][0])

3072

# Convert df_train['Embeddings'] Pandas series to a np.array of float32
X = np.array(df_train['Embeddings'].to_list(), dtype=np.float32)
X.shape

(600, 3072)

Aplicarás el enfoque de t-Distributed Stochastic Neighbor Embedding (t-SNE) para realizar la reducción de dimensionalidad. Esta técnica reduce el número de dimensiones, mientras preserva los clústeres (los puntos que están cerca permanecen cerca). Para los datos originales, el modelo intenta construir una distribución sobre la cual otros puntos de datos son "vecinos" (por ejemplo, comparten un significado similar). Luego optimiza una función objetivo para mantener una distribución similar en la visualización.

tsne = TSNE(random_state=0)
tsne_results = tsne.fit_transform(X)

df_tsne = pd.DataFrame(tsne_results, columns=['TSNE1', 'TSNE2'])
df_tsne['Class Name'] = df_train['Class Name'] # Add labels column from df_train to df_tsne
df_tsne

TSNE1      TSNE2 Class Name
0    40.186756  -0.446112  sci.crypt
1    35.983940   0.442592  sci.crypt
2    29.899532  -9.017360  sci.crypt
3    33.353500   5.954631  sci.crypt
4    24.770006  -8.209906  sci.crypt
..         ...        ...        ...
595 -23.769527  14.458261  sci.space
596 -29.530485  16.145725  sci.space
597 -14.785336  11.442585  sci.space
598 -22.677313  -0.399542  sci.space
599 -15.382431   4.611829  sci.space

[600 rows x 3 columns]

fig, ax = plt.subplots(figsize=(8,6)) # Set figsize
sns.set_style('darkgrid', {"grid.color": ".6", "grid.linestyle": ":"})
sns.scatterplot(data=df_tsne, x='TSNE1', y='TSNE2', hue='Class Name', palette='hls')
sns.move_legend(ax, "upper left", bbox_to_anchor=(1, 1))
plt.title('Scatter plot of news using t-SNE')
plt.xlabel('TSNE1')
plt.ylabel('TSNE2')
plt.axis('equal')

(np.float64(-45.10566120147705),
 np.float64(47.539869117736814),
 np.float64(-41.033332443237306),
 np.float64(37.77369651794434))

<Figure size 800x600 with 1 Axes>

Compara los resultados con KMeans

La agrupación KMeans{:.external} es un algoritmo de agrupación popular y se usa a menudo para el aprendizaje no supervisado. Determina iterativamente los mejores k puntos centrales y asigna cada ejemplo al centroide más cercano. Introduce los embeddings directamente en el algoritmo KMeans para comparar la visualización de los embeddings con el rendimiento de un algoritmo de aprendizaje automático.

# Apply KMeans
kmeans_model = KMeans(n_clusters=4, random_state=1, n_init='auto').fit(X)
labels = kmeans_model.fit_predict(X)

df_tsne['Cluster'] = labels
df_tsne

TSNE1      TSNE2 Class Name  Cluster
0    40.186756  -0.446112  sci.crypt        1
1    35.983940   0.442592  sci.crypt        1
2    29.899532  -9.017360  sci.crypt        1
3    33.353500   5.954631  sci.crypt        1
4    24.770006  -8.209906  sci.crypt        1
..         ...        ...        ...      ...
595 -23.769527  14.458261  sci.space        2
596 -29.530485  16.145725  sci.space        2
597 -14.785336  11.442585  sci.space        2
598 -22.677313  -0.399542  sci.space        2
599 -15.382431   4.611829  sci.space        2

[600 rows x 4 columns]

fig, ax = plt.subplots(figsize=(8,6)) # Set figsize
sns.set_style('darkgrid', {"grid.color": ".6", "grid.linestyle": ":"})
sns.scatterplot(data=df_tsne, x='TSNE1', y='TSNE2', hue='Cluster', palette='magma')
sns.move_legend(ax, "upper left", bbox_to_anchor=(1, 1))
plt.title('Scatter plot of news using KMeans Clustering')
plt.xlabel('TSNE1')
plt.ylabel('TSNE2')
plt.axis('equal')

(np.float64(-45.10566120147705),
 np.float64(47.539869117736814),
 np.float64(-41.033332443237306),
 np.float64(37.77369651794434))

<Figure size 800x600 with 1 Axes>

def get_majority_cluster_per_group(df_tsne_cluster, class_names):
  class_clusters = dict()
  for c in class_names:
    # Get rows of dataframe that are equal to c
    rows = df_tsne_cluster.loc[df_tsne_cluster['Class Name'] == c]
    # Get majority value in Cluster column of the rows selected
    cluster = rows.Cluster.mode().values[0]
    # Populate mapping dictionary
    class_clusters[c] = cluster
  return class_clusters

classes = df_tsne['Class Name'].unique()
class_clusters = get_majority_cluster_per_group(df_tsne, classes)
class_clusters

{'sci.crypt': np.int32(1),
 'sci.electronics': np.int32(0),
 'sci.med': np.int32(3),
 'sci.space': np.int32(2)}

Obtén la mayoría de los clústeres por grupo y mira cuántos de los miembros reales de ese grupo están en ese clúster.

# Convert the Cluster column to use the class name
class_by_id = {v: k for k, v in class_clusters.items()}
df_tsne['Predicted'] = df_tsne['Cluster'].map(class_by_id.__getitem__)

# Filter to the correctly matched rows
correct = df_tsne[df_tsne['Class Name'] == df_tsne['Predicted']]

# Summarise, as a percentage
acc = correct['Class Name'].value_counts() / SAMPLE_SIZE
acc

Class Name
sci.space          0.993333
sci.med            0.986667
sci.crypt          0.973333
sci.electronics    0.960000
Name: count, dtype: float64

# Get predicted values by name
df_tsne['Predicted'] = ''
for idx, rows in df_tsne.iterrows():
  cluster = rows['Cluster']
  # Get key from mapping based on cluster value
  key = list(class_clusters.keys())[list(class_clusters.values()).index(cluster)]
  df_tsne.at[idx, 'Predicted'] = key

df_tsne

TSNE1      TSNE2 Class Name  Cluster  Predicted
0    40.186756  -0.446112  sci.crypt        1  sci.crypt
1    35.983940   0.442592  sci.crypt        1  sci.crypt
2    29.899532  -9.017360  sci.crypt        1  sci.crypt
3    33.353500   5.954631  sci.crypt        1  sci.crypt
4    24.770006  -8.209906  sci.crypt        1  sci.crypt
..         ...        ...        ...      ...        ...
595 -23.769527  14.458261  sci.space        2  sci.space
596 -29.530485  16.145725  sci.space        2  sci.space
597 -14.785336  11.442585  sci.space        2  sci.space
598 -22.677313  -0.399542  sci.space        2  sci.space
599 -15.382431   4.611829  sci.space        2  sci.space

[600 rows x 5 columns]

Para visualizar mejor el rendimiento de KMeans aplicado a tus datos, puedes usar una matriz de confusión. La matriz de confusión te permite evaluar el rendimiento del modelo de clasificación más allá de la precisión. Puedes ver cómo se clasifican los puntos mal clasificados. Necesitarás los valores reales y los valores predichos, que has recopilado en el dataframe anterior.

cm = confusion_matrix(df_tsne['Class Name'].to_list(), df_tsne['Predicted'].to_list())
disp = ConfusionMatrixDisplay(confusion_matrix=cm,
                              display_labels=classes)
disp.plot(xticks_rotation='vertical')
plt.title('Confusion Matrix for Actual and Clustered Newsgroups');
plt.grid(False)

<Figure size 640x480 with 2 Axes>

Próximos pasos

¡Ahora has creado tu propia visualización de embeddings con agrupación! Intenta usar tus propios datos textuales para visualizarlos como embeddings. Puedes realizar una reducción de dimensionalidad para completar el paso de visualización. Ten en cuenta que t-SNE es bueno para agrupar entradas, pero puede tardar más en converger o quedarse atascado en mínimos locales.

También existen otros algoritmos de agrupación además de KMeans, como la agrupación espacial basada en densidad (DBSCAN).

Para aprender a usar otros servicios en la API de Gemini, consulta la guía Primeros pasos.