from datasets import load_dataset
import pandas as pd

def download_data():
    print("Connecting to Hugging Face and downloading the PolEmo2.0 dataset...")
    
    # Load the 'all_text' configuration of the dataset
    dataset = load_dataset("clarin-pl/polemo2-official", "all_text", split="train", trust_remote_code=True)

    # Convert the first 400 samples to a pandas DataFrame
    df = dataset.select(range(400)).to_pandas()

    # Rename the target column to 'label' for pipeline compatibility
    df = df.rename(columns={"target": "label"})
    
    # Drop technical columns, retaining only 'text' and 'label'
    df = df[['text', 'label']]

    # Insert a sequential ID column for tracking purposes
    df.insert(0, 'id', range(1, len(df) + 1))

    # Export the processed dataset to a CSV file
    df.to_csv("dataset_test.csv", index=False, encoding="utf-8")
    print("Success! Saved 400 authentic, diverse reviews to 'dataset_test.csv'.")

if __name__ == "__main__":
    download_data()