dataset_manager.py

import os
import csv
import tempfile
from datetime import datetime
from huggingface_hub import HfApi, hf_hub_download

DATASET_REPO = "buumba641/Zambia-MultiLigual-ASR-Dataset"
METADATA_FILE = "metadata.csv"

def get_api():
    return HfApi(token=os.environ.get("HF_TOKEN"))

def upload_audio(audio_path, language, transcript):
    api = get_api()
    record_id = datetime.now().strftime("%Y%m%d_%H%M%S")
    audio_name = f"{language}_{record_id}.wav"

    # Upload audio file
    api.upload_file(
        path_or_fileobj=audio_path,
        path_in_repo=f"audio/{audio_name}",
        repo_id=DATASET_REPO,
        repo_type="dataset"
    )

    new_record = {
        "id": record_id,
        "file_name": f"audio/{audio_name}",
        "language": language,
        "transcript": transcript,
        "correction": "",
        "rating": "",
        "timestamp": datetime.now().isoformat()
    }

    update_metadata(new_record)
    return record_id

def update_metadata(new_record):
    api = get_api()
    try:
        metadata_path = hf_hub_download(
            repo_id=DATASET_REPO,
            filename=METADATA_FILE,
            repo_type="dataset"
        )
        with open(metadata_path, "r", encoding="utf-8") as f:
            rows = list(csv.DictReader(f))
    except Exception:
        # If the file doesn't exist yet, start fresh
        rows = []

    rows.append(new_record)

    with tempfile.NamedTemporaryFile(mode="w", delete=False, newline="", encoding="utf-8") as tmp:
        writer = csv.DictWriter(tmp, fieldnames=new_record.keys())
        writer.writeheader()
        writer.writerows(rows)
        temp_path = tmp.name

    api.upload_file(
        path_or_fileobj=temp_path,
        path_in_repo=METADATA_FILE,
        repo_id=DATASET_REPO,
        repo_type="dataset",
        commit_message="Add new ASR sample"
    )

def update_feedback(record_id, correction, rating):
    api = get_api()
    
    metadata_path = hf_hub_download(
        repo_id=DATASET_REPO,
        filename=METADATA_FILE,
        repo_type="dataset"
    )

    with open(metadata_path, encoding="utf-8") as f:
        rows = list(csv.DictReader(f))

    # Update the specific record
    for row in rows:
        if row["id"] == record_id:
            row["correction"] = correction
            row["rating"] = rating

    with tempfile.NamedTemporaryFile(mode="w", delete=False, newline="", encoding="utf-8") as tmp:
        writer = csv.DictWriter(tmp, fieldnames=rows[0].keys())
        writer.writeheader()
        writer.writerows(rows)
        temp_path = tmp.name

    api.upload_file(
        path_in_repo=METADATA_FILE,
        path_or_fileobj=temp_path,
        repo_id=DATASET_REPO,
        repo_type="dataset",
        commit_message=f"Update ASR feedback for {record_id}"
    )