app.py

import os
import sys
import subprocess
import gradio as gr
import json
import yaml
import tempfile
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import time
import re
from pathlib import Path

# Set NLTK data directory - use the environment variable or default to a writable location
nltk_data_dir = os.environ.get("NLTK_DATA", "/home/user/local/share/nltk_data")
os.environ["NLTK_DATA"] = nltk_data_dir
os.makedirs(nltk_data_dir, exist_ok=True)
print(f"NLTK data directory set to: {nltk_data_dir}")

# VERSA paths - these are set by the Dockerfile
VERSA_ROOT = "/home/user/app/versa"
VERSA_BIN = os.path.join(VERSA_ROOT, "versa", "bin", "scorer.py")
VERSA_CONFIG_DIR = os.path.join(VERSA_ROOT, "egs", "demo")

# Data directories - also set up by the Dockerfile
DATA_DIR = "/home/user/app/data"
UPLOAD_DIR = os.path.join(DATA_DIR, "uploads")
RESULTS_DIR = os.path.join(DATA_DIR, "results")
CONFIG_DIR = os.path.join(DATA_DIR, "configs")

# Create directories if they don't exist
for directory in [DATA_DIR, UPLOAD_DIR, RESULTS_DIR, CONFIG_DIR]:
    os.makedirs(directory, exist_ok=True)

# Debug information
print(f"Current user ID: {os.getuid()}")
print(f"Current effective user ID: {os.geteuid()}")
print(f"Current working directory: {os.getcwd()}")
print(f"Data directory permissions: {oct(os.stat(DATA_DIR).st_mode)}")
print(f"NLTK data directory permissions: {oct(os.stat(nltk_data_dir).st_mode) if os.path.exists(nltk_data_dir) else 'Not found'}")

# Custom function to handle mixed data types
def format_value(value):
    return f"{value:.2f}" if isinstance(value, float) else value

def convert_python_dict_to_json(py_dict_str):
    result = py_dict_str
    
    # Handle dictionary keys
    result = re.sub(r"'([^']*)'(?=\s*:)", r'"\1"', result)
    
    # Handle string values (this is tricky with nested apostrophes)
    result = re.sub(r':\s*\'([^\']*)\'', lambda m: ': "' + m.group(1).replace('"', '\\"') + '"', result)
    
    # Replace Python's True/False/None with JSON equivalents
    result = result.replace('True', 'true').replace('False', 'false').replace('None', 'null').replace("inf", "Infinity")
    
    return result

# Check if VERSA is installed
def check_versa_installation():
    """Check if VERSA is properly installed"""
    if not os.path.exists(VERSA_ROOT):
        return False, "VERSA directory not found."
    
    if not os.path.exists(VERSA_BIN):
        return False, "VERSA binary not found."
    
    if not os.path.exists(VERSA_CONFIG_DIR):
        return False, "VERSA configuration directory not found."
    
    # Check if the .installation_complete file exists (created by Dockerfile)
    if not os.path.exists(os.path.join(VERSA_ROOT, ".installation_complete")):
        return False, "VERSA installation indicator file not found."
    
    return True, "VERSA is properly installed."

# Check VERSA installation at startup
versa_installed, versa_status = check_versa_installation()
print(f"VERSA installation status: {versa_status}")

# Find available metric configs
def get_available_metrics():
    """Get list of available metrics from VERSA config directory"""
    metrics = []
    
    if not versa_installed:
        # If VERSA is not installed, return an empty list
        return metrics
    
    # Get all YAML files from the egs directory
    for root, _, files in os.walk(VERSA_CONFIG_DIR):
        for file in files:
            if file.endswith('.yaml'):
                path = os.path.join(root, file)
                # Get relative path from VERSA_CONFIG_DIR
                rel_path = os.path.relpath(path, VERSA_CONFIG_DIR)
                metrics.append(rel_path)
    
    # Add custom configs
    for root, _, files in os.walk(CONFIG_DIR):
        for file in files:
            if file.endswith('.yaml'):
                path = os.path.join(root, file)
                rel_path = f"custom/{os.path.basename(path)}"
                metrics.append(rel_path)
    
    return sorted(metrics)

# Get all available metric names
def get_available_metric_names():
    """Get list of all available metric names in VERSA"""
    metric_names = set()
    
    if not versa_installed:
        # If VERSA is not installed, return an empty list
        return []
    
    # First check the universal metrics file
    universal_metrics_yaml = os.path.join(CONFIG_DIR, "universal_metrics.yaml")
    if os.path.exists(universal_metrics_yaml):
        try:
            with open(universal_metrics_yaml, 'r') as f:
                config = yaml.safe_load(f)
                if isinstance(config, list):
                    for item in config:
                        if isinstance(item, dict) and 'name' in item:
                            metric_names.add(item['name'])
        except Exception:
            pass
    
    # Then parse all YAML files to extract additional metric names
    for root, _, files in os.walk(VERSA_CONFIG_DIR):
        for file in files:
            if file.endswith('.yaml'):
                path = os.path.join(root, file)
                try:
                    with open(path, 'r') as f:
                        config = yaml.safe_load(f)
                        if isinstance(config, list):
                            for item in config:
                                if isinstance(item, dict) and 'name' in item:
                                    metric_names.add(item['name'])
                except Exception:
                    pass
    
    return sorted(list(metric_names))

# Get metric description from YAML file
def get_metric_description(metric_path):
    """Get description of a metric from its YAML file"""
    if not versa_installed:
        return "VERSA is not installed. Metric descriptions are unavailable."
    
    if metric_path.startswith("custom/"):
        # Handle custom metrics
        filename = metric_path.split("/")[1]
        full_path = os.path.join(CONFIG_DIR, filename)
    else:
        full_path = os.path.join(VERSA_CONFIG_DIR, metric_path)
    
    try:
        with open(full_path, 'r') as f:
            config = yaml.safe_load(f)
            
            # Check if the config has a description field
            if isinstance(config, dict) and 'description' in config:
                return config.get('description', 'No description available')
            
            # If it's a list of metrics, return a summary
            if isinstance(config, list):
                metric_names = []
                for item in config:
                    if isinstance(item, dict) and 'name' in item:
                        metric_names.append(item['name'])
                
                if metric_names:
                    return f"Contains metrics: {', '.join(metric_names)}"
            
            return "No description available"
    except Exception as e:
        return f"Could not load description: {str(e)}"

# Create custom metric config file
def create_custom_metric_config(selected_metrics, metric_parameters):
    """Create a custom metric configuration file from selected metrics"""
    if not versa_installed:
        return None, "VERSA is not installed. Cannot create custom metric configuration."
    
    if not selected_metrics:
        return None, "Please select at least one metric"
    
    # Load universal metrics as reference
    universal_metrics = []
    universal_metrics_yaml = os.path.join(CONFIG_DIR, "universal_metrics.yaml")
    try:
        with open(universal_metrics_yaml, 'r') as f:
            universal_metrics = yaml.safe_load(f)
    except Exception as e:
        return None, f"Error loading universal metrics: {str(e)}"
    
    # Create new metric config
    custom_metrics = []
    for metric_name in selected_metrics:
        # Find the metric in universal metrics
        for metric in universal_metrics:
            if metric.get('name') == metric_name:
                # Add the metric to custom metrics
                custom_metrics.append(metric.copy())
                break
    
    # Apply any custom parameters from metric_parameters
    if metric_parameters:
        try:
            params = yaml.safe_load(metric_parameters)
            if isinstance(params, dict):
                for metric in custom_metrics:
                    metric_name = metric.get('name')
                    if metric_name in params and isinstance(params[metric_name], dict):
                        # Update metric parameters
                        metric.update(params[metric_name])
        except Exception as e:
            return None, f"Error parsing metric parameters: {str(e)}"
    
    # Create a custom YAML file
    timestamp = int(time.time())
    custom_yaml_path = os.path.join(CONFIG_DIR, f"custom_metrics_{timestamp}.yaml")
    
    try:
        with open(custom_yaml_path, 'w') as f:
            yaml.dump(custom_metrics, f, default_flow_style=False)
        
        return f"custom/{os.path.basename(custom_yaml_path)}", f"Custom metric configuration created successfully with {len(custom_metrics)} metrics"
    except Exception as e:
        return None, f"Error creating custom metric configuration: {str(e)}"

# Load metric config file
def load_metric_config(config_path):
    """Load a metric configuration file and return its content"""
    if not versa_installed and not config_path.startswith("custom/"):
        return "VERSA is not installed. Cannot load metric configuration."
    
    if config_path.startswith("custom/"):
        # Handle custom metrics
        filename = config_path.split("/")[1]
        full_path = os.path.join(CONFIG_DIR, filename)
    else:
        full_path = os.path.join(VERSA_CONFIG_DIR, config_path)
    
    try:
        with open(full_path, 'r') as f:
            content = f.read()
        
        return content
    except Exception as e:
        return f"Error loading metric configuration: {str(e)}"

# Save uploaded YAML file
def save_uploaded_yaml(file_obj):
    """Save an uploaded YAML file to the configs directory"""
    if file_obj is None:
        return None, "No file uploaded"
    
    try:
        # Get the file name and create a new path
        filename = os.path.basename(file_obj.name)
        if not filename.endswith('.yaml'):
            filename += '.yaml'
        
        # Ensure unique filename
        timestamp = int(time.time())
        yaml_path = os.path.join(CONFIG_DIR, f"uploaded_{timestamp}_{filename}")
        
        # Copy the file
        with open(file_obj.name, 'rb') as src, open(yaml_path, 'wb') as dst:
            dst.write(src.read())
        
        # Validate YAML format
        with open(yaml_path, 'r') as f:
            yaml.safe_load(f)
        
        return f"custom/{os.path.basename(yaml_path)}", f"YAML file uploaded successfully as {os.path.basename(yaml_path)}"
    except yaml.YAMLError:
        if os.path.exists(yaml_path):
            os.remove(yaml_path)
        return None, "Invalid YAML format. Please check your file."
    except Exception as e:
        if os.path.exists(yaml_path):
            os.remove(yaml_path)
        return None, f"Error uploading YAML file: {str(e)}"

# Process audio files and run VERSA evaluation
def evaluate_audio(gt_file, pred_file, metric_config, include_timestamps=False):
    """Evaluate audio files using VERSA"""
    if not versa_installed:
        return None, "VERSA is not installed. Evaluation cannot be performed."
    
    if pred_file is None:
        return None, "Please upload the audio file to be evaluated."
    
    # Determine the metric config path
    if metric_config.startswith("custom/"):
        # Handle custom metrics
        filename = metric_config.split("/")[1]
        metric_config_path = os.path.join(CONFIG_DIR, filename)
    else:
        metric_config_path = os.path.join(VERSA_CONFIG_DIR, metric_config)
    
    # Create temp directory for results
    with tempfile.TemporaryDirectory() as temp_dir:
        output_file = os.path.join(temp_dir, "result.json")

        # Create SCP file
        pred_scp_path = os.path.join(temp_dir, "pred.scp")
        with open(pred_scp_path, "w") as pred_scp:
            pred_scp.write("test {}\n".format(pred_file))
       
        # For case without reference audio
        if gt_file is not None:
            gt_scp_path = os.path.join(temp_dir, "gt.scp")
            with open(gt_scp_path, "w") as gt_scp:
                gt_scp.write("test {}\n".format(gt_file))
        else:
            gt_scp_path = "None"

        # Build command
        cmd = [
            sys.executable, VERSA_BIN,
            "--score_config", metric_config_path,
            "--gt", gt_scp_path,
            "--pred", pred_scp_path,
            "--output_file", output_file,
            "--use_gpu", "true",
        ]
        
        if include_timestamps:
            cmd.append("--include_timestamp")
        
        # Run VERSA evaluation
        try:

            # Set environment variables for the subprocess
            env = os.environ.copy()
            env["LIBROSA_CACHE_DIR"] = "/tmp/librosa_cache"
            env["LIBROSA_CACHE_LEVEL"] = "0"
            
            # Pass through the NLTK_DATA environment variable
            env["NLTK_DATA"] = nltk_data_dir

            # Set huggingface cache
            env["HF_HOME"] = "/home/user/.cache/huggingface"

            process = subprocess.run(
                cmd,
                check=True,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                env=env,
            )
            
            # Read results
            if os.path.exists(output_file):
                with open(output_file, 'r') as f:
                    output_file_result = f.read().strip()
                    print(convert_python_dict_to_json(output_file_result))
                    results = json.loads(convert_python_dict_to_json(output_file_result))

                results = {key: format_value(value) for key, value in results.items()}
                # Format results as DataFrame
                if results:
                    results_df = pd.DataFrame([results])
                    return results_df, json.dumps(results, indent=2)
                else:
                    return None, "Evaluation completed but no results were generated."
            else:
                return None, "Evaluation completed but no results file was generated."
        
        except subprocess.CalledProcessError as e:
            return None, f"Error running VERSA: {e.stderr}"

# Create the Gradio interface
def create_gradio_demo():
    """Create the Gradio demo interface"""
    available_metrics = get_available_metrics()
    default_metric = "speech.yaml" if "speech.yaml" in available_metrics else available_metrics[0] if available_metrics else None
    metric_names = get_available_metric_names()
    
    with gr.Blocks(title="VERSA Speech & Audio Evaluation Demo") as demo:
        gr.Markdown("# VERSA: Versatile Evaluation of Speech and Audio")
        
        # Display installation status
        with gr.Row():
            installation_status = gr.Textbox(
                value=f"VERSA Installation Status: {'Installed' if versa_installed else 'Not Installed - ' + versa_status}",
                label="Installation Status",
                interactive=False
            )
        
        if not versa_installed:
            gr.Markdown(f"""
            ## ⚠️ VERSA Not Installed
            
            VERSA does not appear to be properly installed. The Docker container may not have been set up correctly.
            
            Error: {versa_status}
            
            Please check the Docker build logs or contact the administrator.
            """)
        else:
            gr.Markdown("Upload audio files and evaluate them using VERSA metrics.")
            
            with gr.Tabs() as tabs:
                # Standard evaluation tab
                with gr.TabItem("Standard Evaluation"):
                    with gr.Row():
                        with gr.Column():
                            pred_audio = gr.Audio(label="Prediction Audio", type="filepath", sources=["upload", "microphone"])
                            gt_audio = gr.Audio(label="Ground Truth Audio", type="filepath", sources=["upload", "microphone"])
                            
                            metric_dropdown = gr.Dropdown(
                                choices=available_metrics,
                                label="Evaluation Metric Configuration",
                                value=default_metric,
                                info="Select a pre-defined or custom metric configuration"
                            )
                            
                            with gr.Accordion("Metric Configuration Details", open=False):
                                metric_description = gr.Textbox(
                                    label="Metric Description",
                                    value=get_metric_description(default_metric) if default_metric else "",
                                    interactive=False
                                )
                                
                                metric_content = gr.Code(
                                    label="Configuration Content",
                                    language="yaml",
                                    value=load_metric_config(default_metric) if default_metric else "",
                                    interactive=False
                                )
                            
                            # include_timestamps = gr.Checkbox(
                            #     label="Include Timestamps in Results",
                            #     value=False
                            # )
                            
                            eval_button = gr.Button("Evaluate")
                        
                        with gr.Column():
                            results_table = gr.Dataframe(label="Evaluation Results")
                            raw_json = gr.Code(language="json", label="Raw Results")
                
                # Custom metrics creation tab
                with gr.TabItem("Create Custom Metrics"):
                    with gr.Row():
                        with gr.Column():
                            gr.Markdown("### Option 1: Select from Available Metrics")
                            
                            metrics_checklist = gr.CheckboxGroup(
                                choices=metric_names,
                                label="Available Metrics",
                                info="Select the metrics you want to include in your custom configuration"
                            )
                            
                            metric_params = gr.Code(
                                label="Custom Parameters (Optional, YAML format)",
                                language="yaml",
                                interactive=True
                            )
                            
                            create_custom_button = gr.Button("Create Custom Configuration")
                            custom_status = gr.Textbox(label="Status", interactive=False)
                        
                        with gr.Column():
                            gr.Markdown("### Option 2: Upload Your Own YAML File")
                            
                            uploaded_yaml = gr.File(
                                label="Upload YAML Configuration",
                                file_types=[".yaml", ".yml"],
                                type="filepath"
                            )
                            
                            upload_button = gr.Button("Upload Configuration")
                            upload_status = gr.Textbox(label="Upload Status", interactive=False)
                            
                            gr.Markdown("### Generated Configuration")
                            custom_config_path = gr.Textbox(
                                label="Configuration Path",
                                interactive=False,
                                visible=False
                            )
                            
                            custom_config_content = gr.Code(
                                label="Configuration Content",
                                language="yaml",
                                interactive=False
                            )
                
                # About tab
                with gr.TabItem("About VERSA"):
                    gr.Markdown("""
                    ## VERSA: Versatile Evaluation of Speech and Audio
                    
                    VERSA is a toolkit dedicated to collecting evaluation metrics in speech and audio quality. 
                    It provides a comprehensive connection to cutting-edge evaluation techniques and is tightly integrated with ESPnet.
                    
                    With full installation, VERSA offers over 80 metrics with 700+ metric variations based on different configurations. 
                    These metrics encompass evaluations utilizing diverse external resources, including matching and non-matching 
                    reference audio, text transcriptions, and text captions.
                    
                    ### Features
                    
                    - Pythonic interface with flexible configuration
                    - Support for various audio formats and evaluation scenarios
                    - Integration with ESPnet
                    - Batch processing capabilities
                    - Customizable evaluation metrics
                    
                    ### Citation
                    
                    ```
                    @inproceedings{shi2025versa,
                    title={{VERSA}: A Versatile Evaluation Toolkit for Speech, Audio, and Music},
                    author={Jiatong Shi and Hye-jin Shim and Jinchuan Tian and Siddhant Arora and Haibin Wu and Darius Petermann and Jia Qi Yip and You Zhang and Yuxun Tang and Wangyou Zhang and Dareen Safar Alharthi and Yichen Huang and Koichi Saito and Jionghao Han and Yiwen Zhao and Chris Donahue and Shinji Watanabe},
                    booktitle={2025 Annual Conference of the North American Chapter of the Association for Computational Linguistics -- System Demonstration Track},
                    year={2025},
                    url={https://openreview.net/forum?id=zU0hmbnyQm}
                    }
                    
                    @inproceedings{shi2024versaversatileevaluationtoolkit,
                      author={Shi, Jiatong and Tian, Jinchuan and Wu, Yihan and Jung, Jee-Weon and Yip, Jia Qi and Masuyama, Yoshiki and Chen, William and Wu, Yuning and Tang, Yuxun and Baali, Massa and Alharthi, Dareen and Zhang, Dong and Deng, Ruifan and Srivastava, Tejes and Wu, Haibin and Liu, Alexander and Raj, Bhiksha and Jin, Qin and Song, Ruihua and Watanabe, Shinji},
                      booktitle={2024 IEEE Spoken Language Technology Workshop (SLT)}, 
                      title={ESPnet-Codec: Comprehensive Training and Evaluation of Neural Codecs For Audio, Music, and Speech}, 
                      year={2024},
                      pages={562-569},
                      keywords={Training;Measurement;Codecs;Speech coding;Conferences;Focusing;Neural codecs;codec evaluation},
                      doi={10.1109/SLT61566.2024.10832289}
                    }
                    ```
                    
                    Learn more at [VERSA GitHub Repository](https://github.com/shinjiwlab/versa).
                    """)
            
            # Event handlers
            def update_metric_details(metric_path):
                return get_metric_description(metric_path), load_metric_config(metric_path)
            
            metric_dropdown.change(
                fn=update_metric_details,
                inputs=[metric_dropdown],
                outputs=[metric_description, metric_content]
            )
            
            eval_button.click(
                fn=evaluate_audio,
                inputs=[gt_audio, pred_audio, metric_dropdown], # include_timestamps
                outputs=[results_table, raw_json]
            )
            
            # Create custom metric configuration
            def create_and_update_custom_config(selected_metrics, metric_parameters):
                config_path, status = create_custom_metric_config(selected_metrics, metric_parameters)
                if config_path:
                    content = load_metric_config(config_path)
                    # Refresh the available metrics list
                    metrics_list = get_available_metrics()
                    return status, config_path, content, gr.Dropdown.update(choices=metrics_list, value=config_path)
                else:
                    return status, None, "", gr.Dropdown.update(choices=get_available_metrics())
            
            create_custom_button.click(
                fn=create_and_update_custom_config,
                inputs=[metrics_checklist, metric_params],
                outputs=[custom_status, custom_config_path, custom_config_content, metric_dropdown]
            )
            
            # Upload YAML file
            def upload_and_update_yaml(file_obj):
                config_path, status = save_uploaded_yaml(file_obj)
                if config_path:
                    content = load_metric_config(config_path)
                    # Refresh the available metrics list
                    metrics_list = get_available_metrics()
                    return status, config_path, content, gr.Dropdown.update(choices=metrics_list, value=config_path)
                else:
                    return status, None, "", gr.Dropdown.update(choices=get_available_metrics())
            
            upload_button.click(
                fn=upload_and_update_yaml,
                inputs=[uploaded_yaml],
                outputs=[upload_status, custom_config_path, custom_config_content, metric_dropdown]
            )
    
    return demo

# Launch the app
if __name__ == "__main__":
    demo = create_gradio_demo()
    # Use 0.0.0.0 to listen on all interfaces, which is required for Docker
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)