import pandas as pd
from tqdm import tqdm

data_list = []


def get_questions(row):
    global data_list
    random_samples = df.sample(n=num_choices - 1)
    distractors = random_samples["description"].tolist()
    data = {
        "question": "What is " + row["label"] + "?",
        "correct": row["description"],
        "distractors": distractors,
        "knowledge": "{" + row["label"] + " : " + row["description"] + "}",
    }
    data_list.append(data)


debug = False
num_choices = 4

tqdm.pandas(desc="Progress")
df = pd.read_pickle("data/augmented_datasets/pickle/label_description.pkl")
if debug:
    df = df.iloc[:10]

df = df.progress_apply(get_questions, axis=1)

new_df = pd.DataFrame(data_list)
if not debug:
    new_df.to_pickle("data/augmented_datasets/pickle/description_qa_knowledge.pkl")
else:
    __import__("pudb").set_trace()

[((23, 0, 23, 28), 'tqdm.tqdm.pandas', 'tqdm.pandas', (), '', False, 'from tqdm import tqdm\n'), ((24, 5, 24, 75), 'pandas.read_pickle', 'pd.read_pickle', ({(24, 20, 24, 74): '"""data/augmented_datasets/pickle/label_description.pkl"""'}, {}), "('data/augmented_datasets/pickle/label_description.pkl')", True, 'import pandas as pd\n'), ((30, 9, 30, 32), 'pandas.DataFrame', 'pd.DataFrame', ({(30, 22, 30, 31): 'data_list'}, {}), '(data_list)', True, 'import pandas as pd\n')]