Получить максимальное значение из косинусного склеарна и связать его с исходными данными

Question

Я не был программистом, так что это, вероятно, простой вопрос кодирования, не связанный с математикой или наукой о данных. Я не уверен, как лучше всего достичь своей цели, и я хотел бы сделать свой код более pythoni c, поскольку я пытаюсь стать лучше в этом. Итак, вот в чем проблема: у меня есть список предложений, из которых я извлекаю вложения (в виде числовых c массивов). Затем я кластеризовал эти массивы и с помощью KMeans и переставил вложения в список по cluster_id. Затем я вычисляю оценки косинусного сходства между каждым вектором предложения в каждом кластере и вектором соответствующего центра кластера и сохраняю их в другом списке. Если затем я возьму максимальное значение из каждого кластера оценок, я получу оценку для предложения, которое я хочу, и получу в общей сложности 5 оценок: оценки для предложения в каждом из 5 кластеров с максимальной оценкой сходства с cluster_centre.

Что я хочу, тем не менее, это оригинальное предложение назад (которое заносится в список под названием corpus). Я вставил код ниже. Вы заметите, что я получил 5 максимальных баллов, но после этого я понятия не имею, как вернуться к первоначальному предложению. Должен ли я хранить все данные в другой структуре данных, например в словаре или в фрейме данных (и если да, то как)? Есть ли простой способ вернуть исходное предложение? Моя цель - вернуть 5 предложений в конце, максимум для каждого из 5 кластеров (не числовые векторы внедрения, а исходные текстовые предложения).

И менее важный побочный вопрос: Как я могу сделать этот код более pythoni c? Может быть, без использования циклов for, например? Есть ли способ? Спасибо за любую помощь! :)

from sentence_transformers import SentenceTransformer
from sklearn.cluster import AgglomerativeClustering
from sklearn.cluster import KMeans
import numpy as np
import pandas as pd
from sklearn.metrics.pairwise import cosine_similarity

with open('JWN_Nordstrom_MDNA_overview_2017.txt', 'r') as file:
    initial_corpus = file.read()
corpus = initial_corpus.split('. ')

# Extract sentence embeddings
embedder = SentenceTransformer('bert-base-wikipedia-sections-mean-tokens')
corpus_embeddings = embedder.encode(corpus)

# Perform Kmeans clustering
num_clusters = 5
clustering_model = KMeans(n_clusters=num_clusters)
clustering_model.fit(corpus_embeddings)
cluster_assignment = clustering_model.labels_
cluster_centre = clustering_model.cluster_centers_

# Regroup sentence embeddings into 5 clusters
clustered_sentences = [[] for i in range(num_clusters)]
for sentence_id, cluster_id in enumerate(cluster_assignment):
    clustered_sentences[cluster_id].append(corpus_embeddings[sentence_id])

#Calculate cosine similarities
cosine_similarities = [[] for i in range(num_clusters)]
for cluster_id, cluster in enumerate(clustered_sentences):
    for sentence in cluster:
        cosine_similarity_score = cosine_similarity(sentence.reshape(1,-1), cluster_centre[cluster_id].reshape(1,-1))
        cosine_similarities[cluster_id].append(cosine_similarity_score)

for cluster in cosine_similarities:
    print(max(cluster))