annotation/labeler_tool.py

"""GUI Tool for manual annotation of Myanmar speech data."""

import json
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

import pandas as pd

logger = logging.getLogger(__name__)


@dataclass
class AnnotationLabel:
    """Single annotation label."""
    utterance_id: str
    text: str
    sentiment: str
    intensity: float  # 0.0 - 1.0
    confidence: float  # 0.0 - 1.0
    notes: str = ""
    annotator: str = "anonymous"
    timestamp: str = ""


@dataclass
class AnnotationSession:
    """Annotation session data."""
    session_id: str
    samples: List[Dict]
    labels: List[AnnotationLabel] = field(default_factory=list)
    completed_indices: set = field(default_factory=set)
    current_index: int = 0
    
    @property
    def progress(self) -> float:
        if not self.samples:
            return 0.0
        return len(self.completed_indices) / len(self.samples)


class MyanmarAnnotationTool:
    """GUI-based annotation tool for Myanmar speech data."""
    
    # Sentiment classes
    SENTIMENT_CLASSES = [
        ("positive", "အပြုသဘော"),
        ("negative", "အနှုတ်သဘော"),
        ("neutral", "အလယ်အလတ်"),
        ("sarcastic", "သရော်သည်"),
        ("confused", "ရောထွေး"),
    ]
    
    def __init__(self, output_dir: str = "data/annotations/human_annotated"):
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.session: Optional[AnnotationSession] = None
        
    def load_dataset(self, path: str) -> AnnotationSession:
        """Load dataset for annotation."""
        if path.endswith(".jsonl"):
            samples = []
            with open(path, "r", encoding="utf-8") as f:
                for line in f:
                    samples.append(json.loads(line))
        elif path.endswith(".csv"):
            df = pd.read_csv(path)
            samples = df.to_dict("records")
        else:
            raise ValueError(f"Unsupported file format: {path}")
        
        session_id = Path(path).stem
        self.session = AnnotationSession(
            session_id=session_id,
            samples=samples,
        )
        
        logger.info(f"Loaded {len(samples)} samples for annotation")
        return self.session
    
    def get_current_sample(self) -> Optional[Dict]:
        """Get current sample to annotate."""
        if not self.session:
            return None
        if self.session.current_index >= len(self.session.samples):
            return None
        return self.session.samples[self.session.current_index]
    
    def submit_annotation(
        self,
        sentiment: str,
        intensity: float,
        confidence: float,
        notes: str = "",
        annotator: str = "anonymous",
    ) -> bool:
        """Submit an annotation for the current sample."""
        if not self.session:
            logger.error("No active session")
            return False
        
        sample = self.get_current_sample()
        if not sample:
            logger.error("No current sample")
            return False
        
        from datetime import datetime
        
        label = AnnotationLabel(
            utterance_id=sample.get("id", f"utt_{self.session.current_index}"),
            text=sample.get("text", ""),
            sentiment=sentiment,
            intensity=intensity,
            confidence=confidence,
            notes=notes,
            annotator=annotator,
            timestamp=datetime.now().isoformat(),
        )
        
        self.session.labels.append(label)
        self.session.completed_indices.add(self.session.current_index)
        
        # Move to next incomplete sample
        self._advance_to_next()
        
        logger.info(f"Annotated sample {label.utterance_id} as {sentiment}")
        return True
    
    def _advance_to_next(self) -> None:
        """Move to next incomplete sample."""
        if not self.session:
            return
        
        for i in range(self.session.current_index + 1, len(self.session.samples)):
            if i not in self.session.completed_indices:
                self.session.current_index = i
                return
        
        for i in range(self.session.current_index):
            if i not in self.session.completed_indices:
                self.session.current_index = i
                return
    
    def skip_sample(self) -> bool:
        """Skip current sample without annotating."""
        if not self.session:
            return False
        
        self._advance_to_next()
        return True
    
    def go_to_sample(self, index: int) -> bool:
        """Go to specific sample index."""
        if not self.session:
            return False
        if 0 <= index < len(self.session.samples):
            self.session.current_index = index
            return True
        return False
    
    def save_session(self, path: Optional[str] = None) -> str:
        """Save annotation session to file."""
        if not self.session:
            raise ValueError("No active session")
        
        if path is None:
            path = self.output_dir / f"{self.session.session_id}_annotations.jsonl"
        
        labels_data = [
            {
                "utterance_id": label.utterance_id,
                "text": label.text,
                "sentiment": label.sentiment,
                "intensity": label.intensity,
                "confidence": label.confidence,
                "notes": label.notes,
                "annotator": label.annotator,
                "timestamp": label.timestamp,
            }
            for label in self.session.labels
        ]
        
        with open(path, "w", encoding="utf-8") as f:
            for item in labels_data:
                f.write(json.dumps(item, ensure_ascii=False) + "\n")
        
        logger.info(f"Saved {len(labels_data)} annotations to {path}")
        return str(path)
    
    def export_to_dataframe(self) -> pd.DataFrame:
        """Export annotations as DataFrame."""
        if not self.session:
            raise ValueError("No active session")
        
        return pd.DataFrame([
            {
                "utterance_id": label.utterance_id,
                "text": label.text,
                "sentiment": label.sentiment,
                "intensity": label.intensity,
                "confidence": label.confidence,
                "notes": label.notes,
                "annotator": label.annotator,
            }
            for label in self.session.labels
        ])
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get annotation statistics."""
        if not self.session:
            return {}
        
        df = self.export_to_dataframe()
        
        return {
            "total_samples": len(self.session.samples),
            "annotated": len(self.session.labels),
            "remaining": len(self.session.samples) - len(self.session.labels),
            "progress": self.session.progress,
            "sentiment_distribution": df["sentiment"].value_counts().to_dict() if len(df) > 0 else {},
            "annotator_counts": df["annotator"].value_counts().to_dict() if len(df) > 0 else {},
        }


# Gradio-based GUI
def create_gradio_interface(tool: MyanmarAnnotationTool):
    """Create Gradio-based GUI for annotation."""
    import gradio as gr
    
    with gr.Blocks(title="Myanmar Annotation Tool") as app:
        gr.Markdown("# 🇲🇲 Myanmar Speech Annotation Tool")
        
        with gr.Row():
            with gr.Column():
                sample_display = gr.Textbox(
                    label="Text to Annotate",
                    lines=3,
                    interactive=False,
                )
                prosody_display = gr.JSON(label="Prosody Features")
            
            with gr.Column():
                sentiment_dropdown = gr.Dropdown(
                    choices=[s[0] for s in tool.SENTIMENT_CLASSES],
                    label="Sentiment",
                    value="neutral",
                )
                intensity_slider = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.5,
                    step=0.1,
                    label="Intensity",
                )
                confidence_slider = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.8,
                    step=0.1,
                    label="Confidence",
                )
                notes_input = gr.Textbox(
                    label="Notes",
                    lines=2,
                )
        
        with gr.Row():
            submit_btn = gr.Button("Submit", variant="primary")
            skip_btn = gr.Button("Skip")
            save_btn = gr.Button("Save Session")
        
        with gr.Row():
            progress_display = gr.Textbox(label="Progress", interactive=False)
            stats_display = gr.JSON(label="Statistics")
        
        def update_display():
            sample = tool.get_current_sample()
            if sample:
                return (
                    sample.get("text", ""),
                    sample.get("prosody", {}),
                )
            return ("No more samples", {})
        
        def submit_annotation(sentiment, intensity, confidence, notes):
            tool.submit_annotation(sentiment, intensity, confidence, notes)
            sample = tool.get_current_sample()
            stats = tool.get_statistics()
            if sample:
                return (
                    sample.get("text", ""),
                    sample.get("prosody", {}),
                    f"{stats['annotated']}/{stats['total_samples']} ({stats['progress']*100:.1f}%)",
                    stats,
                )
            return ("All samples annotated!", {}, "100%", stats)
        
        def skip():
            tool.skip_sample()
            return update_display()
        
        def save():
            path = tool.save_session()
            return f"Saved to {path}"
        
        submit_btn.click(
            submit_annotation,
            inputs=[sentiment_dropdown, intensity_slider, confidence_slider, notes_input],
            outputs=[sample_display, prosody_display, progress_display, stats_display],
        )
        skip_btn.click(
            skip,
            outputs=[sample_display, prosody_display],
        )
        save_btn.click(
            save,
            outputs=[progress_display],
        )
        
        # Initialize display
        app.load(
            update_display,
            outputs=[sample_display, prosody_display],
        )
    
    return app


if __name__ == "__main__":
    tool = MyanmarAnnotationTool()
    print("MyanmarAnnotationTool initialized")
    print(f"Available sentiment classes: {[s[0] for s in tool.SENTIMENT_CLASSES]}")