sélectionner un nombre variable de tokens dans une colonne pandas en fonction des tuples d'une autre colonne

Question

J'ai un cadre de données avec deux colonnes : sentence contenant du texte et selector contenant des tableaux de tuples de différentes longueurs.

Considérons le cadre de données suivant à titre d'exemple :

df = pd.DataFrame({'sentence': ['KEEP some of the words from this sentence.',
                                'Keep SOME of THE words from this sentence.',
                                'KEEP some OF the WORDS from this sentence.',
                                'Keep SOME of THE words FROM this SENTENCE.'],
                   'selector': [[(10, 0, 1)],
                                [(10, 1, 2), (10, 3, 4)],
                                [(10, 0, 1), (10, 2, 3), (10, 4, 5)],
                                [(10, 1, 2), (10, 3, 4), (10, 5, 6), (10, 7, 8)]]})

Je veux maintenant sélectionner les mots de sentence à la position indiquée par le deuxième élément de chaque tuple (en ignorant l'élément 10 dans chaque tuple). Par exemple, pour la première ligne, je veux que le jeton dans la colonne sentence à la deuxième position de tous les tuples (dont il n'existe qu'un seul : (10, 0, 1) ), c'est-à-dire le jeton à la position 0 : KEEP . (Pour plus de clarté, j'ai écrit tous les mots à sélectionner en MAJUSCULES).

Je voudrais obtenir un cadre de données ressemblant à ceci :

sentence                                    selector                                           selected_tokens
KEEP some of the words from this sentence.  [(10, 0, 1)],                                      ['KEEP']
KEEP some OF the WORDS from this sentence.  [(10, 0, 1), (100, 2, 3), (10, 4, 5)],             ['KEEP', 'OF', 'WORDS']
Keep SOME of THE words from this sentence.  [(10, 1, 2), (10, 3, 4)],                          ['SOME', 'THE']
Keep SOME of THE words FROM this SENTENCE.  [(10, 1, 2), (10, 3, 4), (10, 5, 6), (10, 7, 8)],  ['SOME', 'THE', 'FROM', 'SENTENCE']

L'accès au premier jeton fonctionne bien en utilisant df['tok0_pos'] = df['selector'].str[0].str[1] pour les postes et df['words0'] = [txt.split()[loc] for txt, loc in zip(df['sentence'], df['tok0_pos'])] pour les jetons. Cependant, en raison des longueurs variables (l'ensemble de données réelles contient 0-25 tuples dans la colonne selectors ), cela se bloque rapidement ou est fastidieux.

Quelqu'un peut-il indiquer comment obtenir au mieux la colonne selected_tokens dans l'ensemble de données de l'échantillon ?

Answer 1

Une solution :

df["selected_tokens"] = [[sent[s] for _, s, _ in select] for sent, select in zip(df["sentence"].str.split(), df["selector"])]
print(df["selected_tokens"])

Sortie

0                          [KEEP]
1                     [SOME, THE]
2               [KEEP, OF, WORDS]
3    [SOME, THE, FROM, SENTENCE.]
Name: selected_tokens, dtype: object

Une solution alternative consiste à utiliser numpy pour tirer parti des fonctions d'indexation avancées :

import numpy as np

sentences = df["sentence"].str.split().apply(np.array)
indices = [[s[1] for s in select] for select in df["selector"]]
df["selected_tokens"] = [sentence[i] for sentence, i in zip(sentences, indices)]

Answer 2

Cela devrait fonctionner :

def get_word_at_index(sentence, index):
    return [sentence.split()[i[1]] for i in index]

df.loc[:, 'selected_tokens'] = df.apply(lambda x: get_word_at_index(x["sentence"], x["selector"]), axis=1)