app.py

"""
Memoni Audio Labeling Tool — HuggingFace Spaces (Gradio)
All secrets (HF_TOKEN) stay server-side in Space Secrets.
Audio is served from local cache (loaded from HF parquet server-side).
"""

import os
import json
import soundfile as sf
from datetime import datetime, timezone
from pathlib import Path

import io
import gradio as gr
import pandas as pd
from datasets import load_dataset, Audio
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import EntryNotFoundError

from normalization import REFERENCE_TABLE, normalize, urdu_to_roman

# ── Config ────────────────────────────────────────────────────────────────────
HF_TOKEN    = os.environ.get('HF_TOKEN')
AUDIO_REPO  = 'Aqiba/memoni_clean_audio'
LABELS_REPO = 'Warrior786/memoni_labels'
LABELS_FILE = 'labels.csv'
AUDIO_CACHE = '/tmp/audio_cache'
os.makedirs(AUDIO_CACHE, exist_ok=True)

api = HfApi(token=HF_TOKEN)

# ── Ensure labels dataset repo exists ────────────────────────────────────────
def ensure_labels_repo():
    try:
        api.repo_info(repo_id=LABELS_REPO, repo_type='dataset')
    except Exception:
        api.create_repo(repo_id=LABELS_REPO, repo_type='dataset',
                        private=True, exist_ok=True)
        empty = pd.DataFrame(columns=[
            'volunteer_name', 'video_id', 'transcript', 'script_type',
            'urdu_translation', 'timestamp', 'skipped'
        ])
        path = '/tmp/labels_init.csv'
        empty.to_csv(path, index=False)
        api.upload_file(path_or_fileobj=path, path_in_repo=LABELS_FILE,
                        repo_id=LABELS_REPO, repo_type='dataset',
                        commit_message='Init labels file')

ensure_labels_repo()

# ── Load dataset index (no audio decoding at startup) ────────────────────────
print('Loading dataset index...')
ds_meta = load_dataset(AUDIO_REPO, split='train').remove_columns('audio')
# decode=False → gives raw bytes instead of decoded array — avoids torchcodec entirely
ds_full = load_dataset(AUDIO_REPO, split='train').cast_column('audio', Audio(decode=False))
VIDEO_ID_TO_IDX = {item['video_id']: i for i, item in enumerate(ds_meta)}

MANIFEST = sorted(
    [
        {
            'video_id': item['video_id'],
            'duration': round(item['duration_seconds'], 1),
        }
        for item in ds_meta
    ],
    key=lambda x: x['duration']   # shortest clips first
)
print(f'Manifest ready: {len(MANIFEST)} clips')

# ── Serve audio server-side (cache decoded WAV to /tmp) ──────────────────────
def get_audio_path(video_id: str) -> str:
    path = f'{AUDIO_CACHE}/{video_id}.wav'
    if not os.path.exists(path):
        idx = VIDEO_ID_TO_IDX.get(video_id)
        if idx is None:
            return None
        item = ds_full[idx]
        # decode=False gives {'bytes': b'...', 'path': '...'} — no torchcodec needed
        raw = item['audio']['bytes']
        audio_array, sr = sf.read(io.BytesIO(raw))
        sf.write(path, audio_array, sr)
    return path

# ── Whisper via Groq free tier ────────────────────────────────────────────────
# Groq: 28,800 sec/day free — sign up at console.groq.com, add GROQ_API_KEY to Space Secrets
GROQ_API_KEY = os.environ.get('GROQ_API_KEY')
print(f'[Whisper] GROQ_API_KEY loaded: {bool(GROQ_API_KEY)}')

def _groq_transcribe(audio_path: str, language=None) -> str:
    """Single Groq Whisper call. language=None → auto (outputs Roman/Latin)."""
    from groq import Groq
    client = Groq(api_key=GROQ_API_KEY)
    kwargs = dict(
        model='whisper-large-v3',
        response_format='text',
    )
    if language:
        kwargs['language'] = language
    with open(audio_path, 'rb') as f:
        result = client.audio.transcriptions.create(
            file=(os.path.basename(audio_path), f), **kwargs
        )
    return (result if isinstance(result, str) else result.text).strip()

def get_whisper_drafts(audio_path: str) -> tuple:
    """Returns (roman_draft, urdu_draft). Both '' on failure.
    One Groq call transcribes to Urdu script; transliteration gives Roman."""
    if not audio_path or not os.path.exists(audio_path):
        return '', ''
    if not GROQ_API_KEY:
        print('[Whisper] GROQ_API_KEY not set.')
        return '', ''
    try:
        urdu  = _groq_transcribe(audio_path, language='ur')
        roman = urdu_to_roman(urdu)
        print(f'[Groq Urdu]  {urdu[:60]}')
        print(f'[Groq Roman] {roman[:60]}')
        return roman, urdu
    except Exception as e:
        print(f'[Groq] FAILED: {e}')
        return '', ''


# ── Label storage ─────────────────────────────────────────────────────────────
def load_labels() -> pd.DataFrame:
    try:
        path = hf_hub_download(repo_id=LABELS_REPO, filename=LABELS_FILE,
                               repo_type='dataset', token=HF_TOKEN)
        df = pd.read_csv(path)
        for col in ('urdu_translation', 'whisper_draft', 'whisper_urdu_draft'):
            if col not in df.columns:
                df[col] = ''
        return df
    except EntryNotFoundError:
        return pd.DataFrame(columns=[
            'volunteer_name', 'video_id', 'transcript', 'script_type',
            'urdu_translation', 'whisper_draft', 'whisper_urdu_draft', 'timestamp', 'skipped'
        ])

def get_labeled_ids() -> set:
    return set(load_labels()['video_id'].tolist())

def save_whisper_inference(video_id, whisper_roman, whisper_urdu):
    """Save raw Groq output immediately on clip load — no volunteer action needed."""
    if not whisper_roman and not whisper_urdu:
        return
    df = load_labels()
    # Skip if already have a whisper row for this clip
    if not df[(df['video_id'] == video_id) & (df['volunteer_name'] == '[whisper]')].empty:
        return
    new_row = {
        'volunteer_name':    '[whisper]',
        'video_id':          video_id,
        'transcript':        whisper_roman,
        'script_type':       'roman',
        'urdu_translation':  whisper_urdu,
        'whisper_draft':     whisper_roman,
        'whisper_urdu_draft': whisper_urdu,
        'timestamp':         datetime.now(timezone.utc).isoformat(),
        'skipped':           False,
    }
    df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
    path = '/tmp/labels.csv'
    df.to_csv(path, index=False)
    api.upload_file(
        path_or_fileobj=path,
        path_in_repo=LABELS_FILE,
        repo_id=LABELS_REPO,
        repo_type='dataset',
        token=HF_TOKEN,
        commit_message=f'Whisper inference: {video_id}',
    )
    print(f'[Whisper] Saved inference for {video_id}')

def save_label(volunteer, video_id, transcript, script_type,
               urdu_translation, whisper_draft, whisper_urdu_draft, skipped):
    df = load_labels()
    new_row = {
        'volunteer_name':    volunteer,
        'video_id':          video_id,
        'transcript':        normalize(transcript) if not skipped else '',
        'script_type':       script_type if not skipped else 'skipped',
        'urdu_translation':  urdu_translation.strip() if not skipped else '',
        'whisper_draft':     whisper_draft,       # Roman — always saved
        'whisper_urdu_draft': whisper_urdu_draft, # Urdu  — always saved
        'timestamp':         datetime.now(timezone.utc).isoformat(),
        'skipped':           skipped,
    }
    df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
    path = '/tmp/labels.csv'
    df.to_csv(path, index=False)
    api.upload_file(
        path_or_fileobj=path,
        path_in_repo=LABELS_FILE,
        repo_id=LABELS_REPO,
        repo_type='dataset',
        token=HF_TOKEN,
        commit_message=f'Label: {video_id} by {volunteer}',
    )

# ── Navigation ────────────────────────────────────────────────────────────────
def get_next_clip(labeled_ids: set):
    for item in MANIFEST:
        if item['video_id'] not in labeled_ids:
            return item
    return None

# ── Reference HTML ────────────────────────────────────────────────────────────
def reference_html():
    rows = ''.join(
        f'<tr><td><b>{r}</b></td>'
        f'<td style="font-size:1.1em;padding:2px 8px">{u}</td>'
        f'<td style="color:#888">{e}</td></tr>'
        for r, u, e in REFERENCE_TABLE
    )
    return f"""
    <div style="max-height:500px;overflow-y:auto;font-size:0.85em">
    <table style="width:100%;border-collapse:collapse">
      <thead><tr>
        <th style="text-align:left;padding:4px;border-bottom:1px solid #555">Roman</th>
        <th style="text-align:left;padding:4px;border-bottom:1px solid #555">Urdu</th>
        <th style="text-align:left;padding:4px;border-bottom:1px solid #555">Meaning</th>
      </tr></thead>
      <tbody>{rows}</tbody>
    </table>
    </div>"""

REF_HTML = reference_html()

# ── Gradio UI ─────────────────────────────────────────────────────────────────
with gr.Blocks(title='Memoni Labeling Tool', theme=gr.themes.Soft()) as demo:

    state_volunteer     = gr.State('')
    state_video_id      = gr.State('')
    state_labeled       = gr.State(set())
    state_history       = gr.State([])
    state_whisper_draft = gr.State('')   # Whisper Roman raw output
    state_whisper_urdu  = gr.State('')   # Whisper Urdu raw output

    gr.Markdown('# 🎙️ Memoni Audio Labeling Tool')
    gr.Markdown(
        'Listen to the clip and type what you hear. '
        'Roman, Urdu script, or mixed — your choice. '
        'Urdu translation is optional but very helpful.'
    )

    # ── Name entry ────────────────────────────────────────────────────────────
    with gr.Group(visible=True) as name_group:
        gr.Markdown('### Enter your name to begin')
        name_input  = gr.Textbox(label='Your name', placeholder='e.g. Ahmed')
        name_button = gr.Button('Start Labeling', variant='primary')
        name_status = gr.Markdown('')

    # ── Labeling UI ───────────────────────────────────────────────────────────
    with gr.Group(visible=False) as label_group:
        with gr.Row():
            with gr.Column(scale=3):
                progress_text = gr.Markdown('**Loading...**')
                audio_player  = gr.Audio(label='Audio clip', type='filepath',
                                         interactive=False)
                video_id_box  = gr.Textbox(visible=False)

                whisper_status = gr.Markdown('')
                transcript = gr.Textbox(
                    label='Transcript — what did they say? (Roman / Urdu / Mixed)',
                    placeholder='e.g.  mein ghar aao  or  میں گھر آؤ',
                    lines=2,
                )
                script_type = gr.Radio(
                    choices=['Roman', 'Urdu', 'Mixed'],
                    value='Roman',
                    label='Script used above',
                )
                urdu_translation = gr.Textbox(
                    label='Urdu translation — what does it mean in standard Urdu? (optional)',
                    placeholder='e.g.  میں گھر آؤ  — leave blank if unsure',
                    lines=2,
                )

                with gr.Row():
                    back_btn   = gr.Button('◀ Back', variant='secondary')
                    skip_btn   = gr.Button('⏭ Skip (unclear)', variant='secondary')
                    submit_btn = gr.Button('✅ Submit', variant='primary')
                status_msg = gr.Markdown('')

            with gr.Column(scale=2):
                gr.Markdown('### 📖 Common Words')
                gr.HTML(REF_HTML)

    # ── Helpers ───────────────────────────────────────────────────────────────
    def _load_clip(clip, labeled, history, status_msg):
        audio         = get_audio_path(clip['video_id'])
        roman, urdu   = get_whisper_drafts(audio)
        # Save Groq inference immediately — independent of volunteer action
        save_whisper_inference(clip['video_id'], roman, urdu)
        w_status      = ('🤖 Roman + Urdu drafts — correct if needed'
                         if (roman or urdu) else '✍️ Whisper unavailable — type manually')
        total = len(MANIFEST)
        done  = len(labeled)
        # 11 outputs: audio, vid, progress, msg, labeled, history,
        #             transcript(roman), whisper_status,
        #             state_whisper_draft(roman), state_whisper_urdu(urdu), urdu_translation(urdu)
        return (audio, clip['video_id'],
                f'**Progress: {done} / {total}  |  {clip["duration"]}s**',
                status_msg, labeled, list(history),
                roman, w_status, roman, urdu, urdu)

    # ── Logic ─────────────────────────────────────────────────────────────────
    _EMPTY11 = (None, '', '', '', set(), [], '', '', '', '', '')

    def start(name, _labeled):
        if not name.strip():
            return (gr.update(visible=True), gr.update(visible=False)) + _EMPTY11
        labeled = get_labeled_ids()
        clip    = get_next_clip(labeled)
        total   = len(MANIFEST)
        if clip is None:
            return ((gr.update(visible=False), gr.update(visible=True)) +
                    (None, '', f'🎉 All {total} clips labeled!', '', labeled, [], '', '', '', '', ''))
        vals = _load_clip(clip, labeled, [], '')
        return (gr.update(visible=False), gr.update(visible=True)) + vals

    def submit(volunteer, video_id, text, script, urdu_trans,
               labeled, history, whisper_roman, whisper_urdu):
        if not text.strip():
            return (None, video_id, '⚠️ Transcript is empty — write something or Skip.',
                    '', labeled, history, '', '✍️ Please type a transcript',
                    whisper_roman, whisper_urdu, urdu_trans)
        labeled = set(labeled)
        history = list(history) + [video_id]
        save_label(volunteer, video_id, text, script, urdu_trans,
                   whisper_roman, whisper_urdu, skipped=False)
        labeled.add(video_id)
        clip = get_next_clip(labeled)
        if clip is None:
            total = len(MANIFEST)
            return (None, '', f'**Progress: {total}/{total}**',
                    '🎉 All done!', labeled, history, '', '', '', '', '')
        return _load_clip(clip, labeled, history, '✅ Saved!')

    def skip(volunteer, video_id, labeled, history, whisper_roman, whisper_urdu):
        labeled = set(labeled)
        history = list(history) + [video_id]
        save_label(volunteer, video_id, '', '', '', whisper_roman, whisper_urdu, skipped=True)
        labeled.add(video_id)
        clip = get_next_clip(labeled)
        if clip is None:
            total = len(MANIFEST)
            return (None, '', f'**Progress: {total}/{total}**',
                    '🎉 All done!', labeled, history, '', '', '', '', '')
        return _load_clip(clip, labeled, history, '⏭ Skipped.')

    def go_back(current_video_id, labeled, history):
        history = list(history)
        if not history:
            audio = get_audio_path(current_video_id) if current_video_id else None
            return (audio, current_video_id, '', '⚠️ Already at first clip.',
                    labeled, history, '', '', '', '', '')
        prev_video_id = history[-1]
        history       = history[:-1]

        df       = load_labels()
        prev_row = df[df['video_id'] == prev_video_id]
        prev_transcript = ''
        if not prev_row.empty:
            prev_transcript = str(prev_row.iloc[-1]['transcript'])
            df = df.drop(prev_row.index[-1])
            path = '/tmp/labels.csv'
            df.to_csv(path, index=False)
            api.upload_file(path_or_fileobj=path, path_in_repo=LABELS_FILE,
                            repo_id=LABELS_REPO, repo_type='dataset', token=HF_TOKEN,
                            commit_message=f'Undo: {prev_video_id}')

        labeled = set(labeled)
        labeled.discard(prev_video_id)
        audio     = get_audio_path(prev_video_id)
        total     = len(MANIFEST)
        clip_info = next((c for c in MANIFEST if c['video_id'] == prev_video_id), {})
        dur       = clip_info.get('duration', '')
        roman, urdu = get_whisper_drafts(audio)
        w_st = '📝 Previous label restored — edit and resubmit' if prev_transcript else '✍️ Type transcript'
        return (audio, prev_video_id,
                f'**Progress: {len(labeled)} / {total}  |  {dur}s**',
                '◀ Back — edit and resubmit',
                labeled, history, prev_transcript, w_st, roman, urdu, urdu)

    # ── Events ────────────────────────────────────────────────────────────────
    # 11 main outputs: audio, vid, progress, status, labeled, history,
    #                  transcript, whisper_status, state_whisper_draft,
    #                  state_whisper_urdu, urdu_translation
    _main_outputs = [audio_player, video_id_box, progress_text, status_msg,
                     state_labeled, state_history,
                     transcript, whisper_status, state_whisper_draft,
                     state_whisper_urdu, urdu_translation]

    name_button.click(lambda n: n, inputs=[name_input], outputs=[state_volunteer])
    name_button.click(lambda: '⏳ Loading clips — first load takes ~30 seconds...',
                      outputs=[name_status])
    name_button.click(
        start,
        inputs=[name_input, state_labeled],
        outputs=[name_group, label_group] + _main_outputs,
    )
    submit_btn.click(
        submit,
        inputs=[state_volunteer, video_id_box, transcript, script_type,
                urdu_translation, state_labeled, state_history,
                state_whisper_draft, state_whisper_urdu],
        outputs=_main_outputs,
    )
    skip_btn.click(
        skip,
        inputs=[state_volunteer, video_id_box, state_labeled,
                state_history, state_whisper_draft, state_whisper_urdu],
        outputs=_main_outputs,
    )
    back_btn.click(
        go_back,
        inputs=[video_id_box, state_labeled, state_history],
        outputs=_main_outputs,
    )


if __name__ == '__main__':
    demo.launch()