Разбить словарь, в котором значения ключей - это несколько списков, на поезд и набор тестов Python

Question

Скажем, у меня есть словарь с двумя ключами, спамом и ветчиной, для текста спама и ветчины или сообщений электронной почты, который выглядит следующим образом:

data = {
    'spam': [
        ['hi', "what's", 'going', 'on', 'sexy', 'thing'], 
        ['1-800', 'call', 'girls', 'if', "you're", 'lonely'], 
        ['sexy', 'girls', 'for', 'youuuuuu']], 
    'ham': [['hey', 'hey', 'I', 'got', 'your', 'message,', "I'll", 'be', 'home', 'soon!!!'], 
        ['Madden', 'MUT', 'time', 'boys']]}

Я хочу разбить словарь на обучающие и тестовые наборы (начиная с 80/20 тренировок для тестирования). Я хочу, чтобы разделение было независимо от ключа, поэтому просто задайте 80% от общего количества сообщений для моего обучающего набора и 20% от общего количества сообщений для моего тестового набора. В этом небольшом примере у нас всего 5 сообщений (3 в спаме и 2 в ветчине). Я искал решения, но пока не нашел ничего, что могло бы справиться с такой ситуацией.

Answer 1

Использование метко названного sklearn.model_selection.train_test_split():

from sklearn.model_selection import train_test_split

data = {
    'spam': [
        ['hi', "what's", 'going', 'on', 'sexy', 'thing'],
        ['1-800', 'call', 'girls', 'if', "you're", 'lonely'],
        ['sexy', 'girls', 'for', 'youuuuuu']],
    'ham': [['hey', 'hey', 'I', 'got', 'your', 'message,', "I'll", 'be', 'home', 'soon!!!'],
            ['Madden', 'MUT', 'time', 'boys']]}

all_messages = [(words, k) for k, v in data.items() for words in v]

train, test = train_test_split(list(all_messages), test_size=0.2)

---

Вы можете и, вероятно, должны использовать что-то более мощное, например Pandas:

import pandas as pd
from sklearn.model_selection import train_test_split

data_dict = {
    'spam': [
        ['hi', "what's", 'going', 'on', 'sexy', 'thing'],
        ['1-800', 'call', 'girls', 'if', "you're", 'lonely'],
        ['sexy', 'girls', 'for', 'youuuuuu']],
    'ham': [['hey', 'hey', 'I', 'got', 'your', 'message,', "I'll", 'be', 'home', 'soon!!!'],
            ['Madden', 'MUT', 'time', 'boys']]}

df = pd.DataFrame(data=((k, words) for k, v in data_dict.items() for words in v))

print(df)

train, test = train_test_split(df, test_size=0.2)

print(train)
print(test)

Выход:

      0                                                  1
0  spam               [hi, what's, going, on, sexy, thing]
1  spam           [1-800, call, girls, if, you're, lonely]
2  spam                       [sexy, girls, for, youuuuuu]
3   ham  [hey, hey, I, got, your, message,, I'll, be, h...
4   ham                          [Madden, MUT, time, boys]

      0                                                  1
1  spam           [1-800, call, girls, if, you're, lonely]
2  spam                       [sexy, girls, for, youuuuuu]
0  spam               [hi, what's, going, on, sexy, thing]
3   ham  [hey, hey, I, got, your, message,, I'll, be, h...

     0                          1
4  ham  [Madden, MUT, time, boys]

Answer 2

Вы можете преобразовать dict в список кортежей и затем выполнить разбиение.

>>> l = [(sentence, k) for k,v in data.items() for sentence in v]
>>> random.shuffle(l)
>>> train_size = int(len(l)*0.8)
>>> train, test = l[:train_size], l[train_size:]
>>> len(train)
4
>>> len(test)
1

Каждый элемент представляет собой пару (sentence, label)

>>> test[0]
(['Madden', 'MUT', 'time', 'boys'], 'ham')