storage.py

import os
import json
import threading
from datetime import datetime, timezone

FEEDBACK_FILE = "feedback_data.jsonl"
DATASET_REPO_ID = os.getenv("DATASET_REPO_ID", "")
PUSH_EVERY_N_EVENTS = int(os.getenv("PUSH_EVERY_N_EVENTS", "25"))

_write_lock = threading.Lock()


def now_iso():
    return datetime.now(timezone.utc).isoformat()


def save_record(record):
    with _write_lock:
        with open(FEEDBACK_FILE, "a", encoding="utf-8") as f:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")


def count_records():
    if not os.path.exists(FEEDBACK_FILE):
        return 0
    with open(FEEDBACK_FILE, "r", encoding="utf-8") as f:
        return sum(1 for _ in f)


def export_jsonl_to_dataset_dir(model_name):
    os.makedirs("dataset_export", exist_ok=True)

    train_path = os.path.join("dataset_export", "train.jsonl")
    if os.path.exists(FEEDBACK_FILE):
        with open(FEEDBACK_FILE, "r", encoding="utf-8") as src, open(train_path, "w", encoding="utf-8") as dst:
            dst.write(src.read())

    readme = f"""---
license: mit
task_categories:
- text-generation
- text-classification
language:
- ar
- en
pretty_name: Nanochat Moroccan Feedback Dataset
---

# Nanochat Moroccan Feedback Dataset

This dataset contains:
- conversation turns
- model generations
- likes/dislikes
- unsafe-output flags
- bad-output flags

Source model: `{model_name}`
"""
    with open(os.path.join("dataset_export", "README.md"), "w", encoding="utf-8") as f:
        f.write(readme)


def push_dataset_to_hub(model_name):
    if not DATASET_REPO_ID:
        return False, "DATASET_REPO_ID is not set."

    try:
        from huggingface_hub import create_repo, upload_folder

        export_jsonl_to_dataset_dir(model_name)

        create_repo(
            repo_id=DATASET_REPO_ID,
            repo_type="dataset",
            exist_ok=True,
        )

        upload_folder(
            repo_id=DATASET_REPO_ID,
            repo_type="dataset",
            folder_path="dataset_export",
            path_in_repo=".",
        )
        return True, f"Pushed to {DATASET_REPO_ID}"
    except Exception as e:
        return False, str(e)


def maybe_periodic_push(model_name):
    if not DATASET_REPO_ID:
        return None

    n = count_records()
    if n == 0:
        return None

    if n % PUSH_EVERY_N_EVENTS == 0:
        ok, message = push_dataset_to_hub(model_name)
        return {
            "ok": ok,
            "message": message,
            "record_count": n,
        }

    return None