scripts/export_dataset.py

"""Export Gen_AI Dataset.xlsx into train.csv and test.csv at repo root."""
from pathlib import Path
import csv
import openpyxl

ROOT = Path(__file__).resolve().parent.parent
SRC = ROOT / "Gen_AI Dataset.xlsx"


def export_sheet(ws, out_path: Path) -> int:
    rows = list(ws.iter_rows(values_only=True))
    with out_path.open("w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        for row in rows:
            w.writerow(["" if c is None else c for c in row])
    return len(rows) - 1


def main() -> None:
    wb = openpyxl.load_workbook(SRC, data_only=True)
    train_n = export_sheet(wb["Train-Set"], ROOT / "train.csv")
    test_n = export_sheet(wb["Test-Set"], ROOT / "test.csv")
    print(f"train.csv: {train_n} labeled rows")
    print(f"test.csv:  {test_n} queries")


if __name__ == "__main__":
    main()