Repository: LAION-AI/Open-Assistant
File: data/datasets/poetry_instruction/prepare.py (Python)

Code before:
import json
import os
import random

import pandas as pd

# Read the CSV file into a pandas DataFrame. The CSV file was downloaded from Kaggle, and the source is linked in the README.

CURRENT_DIR = os.path.dirname(__file__)  # Gets directory path of the current python module
df = pd.read_csv(os.path.join(CURRENT_DIR, "PoetryFoundationData.csv"))

# The data in the CSV file is not formatted correctly, so we need to clean it up.
df["Title"] = df["Title"].replace("\n", "", regex=True).replace("\r", "", regex=True)
df["Title"] = df["Title"].str.strip()
df["Title"] = df["Title"].apply(lambda x: f'"{x}"')
df["Poem"] = df["Poem"].str.strip()
df["Poem"] = df["Poem"].str.replace("Translated from the French", "")

# "writing_prompts" are for tasks requesting the assistant to write a poem.
# "topic" or "notTopic" are used depending if the original dataset had a topic listed for the poem or not.
writing_prompts_topic = [
    "Write me a poem about $topic.",
    "I want a poem about $topic.",
    "Can you write a poem? Make it about $topic.",
    "Compose a poem, about $topic.",
    "Make a poem with themes of $topic." "Generate a poem with the following themes: $topic.",
]

writing_prompts_notTopic = [
    "Write me a poem.",
    "I want a poem.",
    "Can you write a poem?",
    "Compose a poem.",
    "Make a poem.",
    "Generate a poem.",
]

# These are replies that the assistant can give to the user.
replies_topic = [
    "Here's a poem about $topic: \n$title\n$poem",
    "Sure, I can do that. Here's a poem about $topic. I call it $title: \n$poem",
    "Okay, a poem about $topic: \n$title\n$poem",
    "Of course! It's called $title: \n$poem",
    "It's called $title: \n$poem",
    "Here's your poem about $topic: \n$title\n$poem",
    "I've written a poem for you about $topic. The title is $title: \n$poem",

Review comment (suggestion):
you can also do something like
`df["Title"].apply(lambda x: "".join(x.splitlines()).strip())`
but it's fine

Code after:
import json
import os
import random

import kaggle
import pandas as pd

# Authenticate the Kaggle API client
kaggle.api.authenticate()

# Download and extract the dataset to the download_path directory
download_path = os.path.join(os.getcwd(), "data", "datasets", "poetry_instruction")
kaggle.api.dataset_download_files("tgdivy/poetry-foundation-poems", path=download_path, unzip=True)

# Read the CSV file into a pandas dataframe
csv_file = os.path.join(download_path, "PoetryFoundationData.csv")
df = pd.read_csv(csv_file)

# The data in the CSV file is not formatted correctly, so we need to clean it up.
df["Title"] = df["Title"].replace("\n", "", regex=True).replace("\r", "", regex=True)
df["Title"] = df["Title"].str.strip()
df["Title"] = df["Title"].apply(lambda x: f'"{x}"')
df["Poem"] = df["Poem"].str.strip()
df["Poem"] = df["Poem"].str.replace("Translated from the French", "")

# "writing_prompts" are for tasks requesting the assistant to write a poem.
# "topic" or "notTopic" are used depending if the original dataset had a topic listed for the poem or not.
writing_prompts_topic = [
    "Write me a poem about $topic.",
    "I want a poem about $topic.",
    "Can you write a poem? Make it about $topic.",
    "Compose a poem, about $topic.",
    "Make a poem with themes of $topic." "Generate a poem with the following themes: $topic.",
]

writing_prompts_notTopic = [
    "Write me a poem.",
    "I want a poem.",
    "Can you write a poem?",
    "Compose a poem.",
    "Make a poem.",
    "Generate a poem.",
]

# These are replies that the assistant can give to the user.
replies_topic = [