diff --git a/.gitignore b/.gitignore index 2b51c6e..b87dd4a 100644 --- a/.gitignore +++ b/.gitignore @@ -109,4 +109,6 @@ climateguard/experiments/claim_extraction/climateguard_claim_extraction .aws/ germany_s3_bucket/ france_s3_bucket/ -.gcloud \ No newline at end of file +.gcloud +.claude +CLAUDE.md \ No newline at end of file diff --git a/data/update_huggingface_dataset.py b/data/update_huggingface_dataset.py new file mode 100644 index 0000000..a899322 --- /dev/null +++ b/data/update_huggingface_dataset.py @@ -0,0 +1,137 @@ +import os +from datetime import date, datetime, timedelta +from typing import List + +import numpy as np +import pandas as pd +import psycopg +import requests +from datasets import Dataset, DatasetDict, load_dataset +from dotenv import load_dotenv + +load_dotenv(os.path.join(os.path.dirname(__file__), ".env")) + + +def fetch_articles(columns: List[str] = None) -> pd.DataFrame: + conninfo = ( + f"host={os.getenv('PG_HOST', 'localhost')} port={os.getenv('PG_PORT', 5432)} " + f"dbname={os.getenv('PG_DATABASE', 'postgres')} user={os.getenv('PG_USER', 'user')} " + f"password={os.getenv('PG_PASSWORD', 'supersecret')}" + ) + columns_str = ", ".join(columns) if columns else "*" + query = f""" + SELECT {columns_str} + FROM analytics.task_global_completion + where task_completion_aggregate_id is not null + """ + with psycopg.connect(conninfo) as conn: + with conn.cursor() as cur: + cur.execute(query) + rows = cur.fetchall() + columns = [desc.name for desc in cur.description] + + return pd.DataFrame(rows, columns=columns) + + +def get_week_number(record): + format_str = "%Y-%m-%dT%H:%M:%S" + if len(record["data_item_start"]) > 19: + format_str = format_str + "%z" + week_number = ( + datetime.strptime(record["data_item_start"], format_str).isocalendar().week + ) + return week_number + + +def split_df_by_week(df): + df["week_number"] = df.data_item_start.dt.isocalendar().week + df_train = df.loc[df.week_number % 5 > 0] + df_test = df.loc[df.week_number % 5 == 0] + return df_train, df_test + +def cast_to_int_with_nan(col: pd.Series, int_type: str='uint8'): + col = col.fillna(-1) + col = col.astype(int_type) + col[col==-1] = np.nan + return col + + +def format_dtypes(df: pd.DataFrame): + df.data_item_day = df.data_item_day.astype('f2') + df.data_item_month = df.data_item_month.astype('f2') + df.data_item_year = df.data_item_year.astype('f2') + df.data_item_model_result = df.data_item_model_result.astype('f2') + df.mesinfo_correct = df.mesinfo_correct.astype('f2') + df.mesinfo_incorrect = df.mesinfo_incorrect.astype('f2') + df.speaker_journalist = df.speaker_journalist.astype('f2') + df.speaker_commentator = df.speaker_commentator.astype('f2') + df.speaker_guest = df.speaker_guest.astype('f2') + df.speaker_politician = df.speaker_politician.astype('f2') + df.speaker_audience = df.speaker_audience.astype('f2') + df.speaker_unknown = df.speaker_unknown.astype('f2') + df.mesinfo_corrected_bool = df.mesinfo_corrected_bool.astype('f2') + df.mesinfo_corrected = df.mesinfo_corrected.fillna("") + return df + +def generate_hf_dataset(df_train, df_test): + df_train.data_item_start = df_train.data_item_start.dt.strftime( + "%Y-%m-%dT%H:%M:%S%z" + ) + df_test.data_item_start = df_test.data_item_start.dt.strftime("%Y-%m-%dT%H:%M:%S%z") + dataset = DatasetDict( + { + "train": Dataset.from_pandas(df_train), + "test": Dataset.from_pandas(df_test), + } + ) + return dataset + + +DATASET_COLUMNS = [ + "task_completion_aggregate_id", + "task_aggregate_id", + "created_at", + "updated_at", + "is_labeled", + "project_id", + "country", + "data_item_id", + "data_item_channel", + "data_item_channel_name", + "data_item_channel_title", + "data_item_channel_program", + "data_item_channel_program_type", + "data_item_day", + "data_item_month", + "data_item_year", + "data_item_start", + "data_item_model_name", + "data_item_model_reason", + "data_item_model_result", + "data_item_plaintext", + "data_item_plaintext_whisper", + "data_item_url_mediatree", + "mesinfo_choice", + "locuteur_choice", + "mesinfo_correct", + "mesinfo_incorrect", + "speaker_journalist", + "speaker_commentator", + "speaker_guest", + "speaker_politician", + "speaker_audience", + "speaker_unknown", + "mesinfo_corrected", + "mesinfo_corrected_bool", + "debunk_references", + "claims", + "explanations", + "other_comments", +] + +if __name__ == "__main__": + df = fetch_articles(DATASET_COLUMNS) + df = format_dtypes(df) + df_train, df_test = split_df_by_week(df) + dataset = generate_hf_dataset(df_train, df_test) + dataset.push_to_hub("DataForGood/climateguard-training", private=True, token=True)