app.py

import os
import gc
import gradio as gr
import torch
import numpy as np
import librosa
import pandas as pd
from transformers import WhisperProcessor, AutoConfig, AutoModel, WhisperConfig, WhisperPreTrainedModel
from transformers.models.whisper.modeling_whisper import WhisperEncoder
import torch.nn as nn
import psutil
import json

i18n = gr.I18n(
    en={
        "subtitle": "Identify any of the 68 languages of Mexico and their 360 variants",
        "input_audio_header": "### 🎙️ 1. Input Audio",
        "upload_or_record_label": "Upload or Record",
        "advanced_options": "⚙️ Advanced Options",
        "language_family": "#### Language Family",
        "top_k_label": "Top-K",
        "threshold_label": "Threshold",
        "superlanguage": "#### Superlanguage",
        "language_code": "#### Language Code",
        "clear_btn": "🗑️ Clear",
        "classify_btn": "🚀 Classify",
        "results_header": "### 📊 2. Classification Results",
        "pred_fam_label": "🌍 Language Family",
        "pred_super_label": "🗣️ Superlanguage",
        "pred_code_label": "🔤 Language Code",
        "curated_header": "### 📋 Curated Samples for Reference",
        "curated_desc": "Select any of the sample results below to load the corresponding audio for inference.",
        "prediction_header": "Prediction",
        "confidence_header": "Confidence",
        "col_sample": "🔊 Sample",
        "col_pred_fam": "Pred Family",
        "col_pred_super": "Pred Super",
        "col_pred_code": "Pred Code",
        "col_act_fam": "Actual Family",
        "col_act_super": "Actual Super",
        "col_act_code": "Actual Code",
        "col_match": "Match"
    },
    es={
        "subtitle": "Identifica cualquiera de las 68 lenguas originarias de México y sus 360 variantes",
        "input_audio_header": "### 🎙️ 1. Entrada de Audio",
        "upload_or_record_label": "Subir o Grabar",
        "advanced_options": "⚙️ Opciones Avanzadas",
        "language_family": "#### Familia Lingüística",
        "top_k_label": "Top-K",
        "threshold_label": "Umbral",
        "superlanguage": "#### Superlengua",
        "language_code": "#### Código de Lengua",
        "clear_btn": "🗑️ Limpiar",
        "classify_btn": "🚀 Clasificar",
        "results_header": "### 📊 2. Resultados de Clasificación",
        "pred_fam_label": "🌍 Familia Lingüística",
        "pred_super_label": "🗣️ Superlengua",
        "pred_code_label": "🔤 Código de Lengua",
        "curated_header": "### 📋 Muestras Curadas de Referencia",
        "curated_desc": "Seleccione cualquiera de los resultados de muestra a continuación para cargar el audio correspondiente para la clasificación.",
        "prediction_header": "Predicción",
        "confidence_header": "Confianza",
        "col_sample": "🔊 Muestra",
        "col_pred_fam": "Familia Pred",
        "col_pred_super": "Super Pred",
        "col_pred_code": "Cód Pred",
        "col_act_fam": "Familia Real",
        "col_act_super": "Super Real",
        "col_act_code": "Cód Real",
        "col_match": "Coincidencia"
    }
)

# --- CONFIGURATION ---
MAX_AUDIO_SECONDS = 30

torch.set_num_threads(1)

# === CUSTOM MODEL CLASSES ===
class WhisperEncoderOnlyConfig(WhisperConfig):
    model_type = "whisper_encoder_classifier"
    
    def __init__(self, n_fam=None, n_super=None, n_code=None, **kwargs):
        super().__init__(**kwargs)
        self.n_fam = n_fam
        self.n_super = n_super
        self.n_code = n_code

class WhisperEncoderOnlyForClassification(WhisperPreTrainedModel):
    config_class = WhisperEncoderOnlyConfig

    def __init__(self, config):
        super().__init__(config)
        self.encoder = WhisperEncoder(config)
        hidden = config.d_model
        self.fam_head = nn.Linear(hidden, config.n_fam)
        self.super_head = nn.Linear(hidden, config.n_super)
        self.code_head = nn.Linear(hidden, config.n_code)
        self.post_init()

    def get_input_embeddings(self):
        return None

    def set_input_embeddings(self, value):
        pass

    def enable_input_require_grads(self):
        return

    def forward(self, input_features, labels=None):
        enc_out = self.encoder(input_features=input_features)
        pooled = enc_out.last_hidden_state.mean(dim=1)

        fam_logits = self.fam_head(pooled)
        super_logits = self.super_head(pooled)
        code_logits = self.code_head(pooled)

        loss = None
        if labels is not None:
            fam_labels, super_labels, code_labels = labels
            loss_fn = nn.CrossEntropyLoss()
            loss = (
                loss_fn(fam_logits, fam_labels) +
                loss_fn(super_logits, super_labels) +
                loss_fn(code_logits, code_labels)
            )

        return {
            "loss": loss,
            "fam_logits": fam_logits,
            "super_logits": super_logits,
            "code_logits": code_logits,
        }

class LabelExtractor:
    """
    Extracts family/super/code labels from tokenized sequences based on training design.
    """
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer

        self.family_tokens = []
        self.super_tokens = []
        self.code_tokens = []

        # Extract special tokens that represent categories from added_vocab
        for token_str, token_id in tokenizer.get_added_vocab().items():
            if token_str.startswith("<|") and token_str.endswith("|>"):
                if token_str in ["<|startoftranscript|>", "<|endoftext|>",
                                "<|nospeech|>", "<|notimestamps|>"]:
                    continue

                if token_str.startswith("<|@"):
                    self.family_tokens.append((token_str, token_id))
                elif self._is_super_token(token_str):
                    self.super_tokens.append((token_str, token_id))
                else:
                    self.code_tokens.append((token_str, token_id))

        # Sort by token_id to match model indices
        self.family_tokens.sort(key=lambda x: x[1])
        self.super_tokens.sort(key=lambda x: x[1])
        self.code_tokens.sort(key=lambda x: x[1])
        
        # We only need the flat lists of token names for inference mapping
        self.family_labels = [tok for tok, _ in self.family_tokens]
        self.super_labels = [tok for tok, _ in self.super_tokens]
        self.code_labels = [tok for tok, _ in self.code_tokens]

        print(f"Extracted labels:")
        print(f"  Families: {len(self.family_labels)}")
        print(f"  Superlanguages: {len(self.super_labels)}")
        print(f"  Codes: {len(self.code_labels)}")

    def _is_super_token(self, token_str):
        # Based on training heuristic
        return len(token_str) > 2 and token_str[2].isupper() and not token_str.startswith("<|@")

# === REGISTER MODEL ===
AutoConfig.register("whisper_encoder_classifier", WhisperEncoderOnlyConfig)
AutoModel.register(WhisperEncoderOnlyConfig, WhisperEncoderOnlyForClassification)

# === LOAD MODEL ===
MODEL_REPO = "tachiwin/language_classification_enconly_model_2"

print("Loading model on CPU...")
processor = WhisperProcessor.from_pretrained(MODEL_REPO)
model = WhisperEncoderOnlyForClassification.from_pretrained(
    MODEL_REPO,
    low_cpu_mem_usage=True
)
model.eval()

# Initialize LabelExtractor to build text mappings
label_extractor = LabelExtractor(processor.tokenizer)

# Load languages mapping
print("Loading language mappings...")
try:
    with open("languages.json", "r", encoding="utf-8") as f:
        languages_data = json.load(f)
    CODE_TO_NAME = {item.get("code"): item.get("inali_name") for item in languages_data if item.get("code") and item.get("inali_name")}
except Exception as e:
    print(f"Warning: Could not load languages.json: {e}")
    CODE_TO_NAME = {}

print("Model loaded successfully!")

def get_mem_usage():
    process = psutil.Process(os.getpid())
    return process.memory_info().rss / (1024 ** 2)

# === INFERENCE FUNCTION ===
def predict_language(audio_path, fam_k=1, fam_thresh=0.0, super_k=1, super_thresh=0.0, code_k=3, code_thresh=0.0):
    if not audio_path:
        raise gr.Error("No audio provided! Please upload or record an audio file.")
    
    gc.collect() 
    start_mem = get_mem_usage()
    print(f"\n--- [LOG] New Request ---")
    print(f"[LOG] Start Memory: {start_mem:.2f} MB")
    
    try:
        print("[LOG] Step 1: Loading and resampling audio from file...")
        audio_array, sample_rate = librosa.load(audio_path, sr=16000)
        
        audio_len_sec = len(audio_array) / 16000
        print(f"[LOG] Audio duration: {audio_len_sec:.2f}s, SR: 16000")
        print(f"[LOG] Memory after load: {get_mem_usage():.2f} MB")
        
        if audio_len_sec > MAX_AUDIO_SECONDS:
            del audio_array
            gc.collect()
            raise gr.Error(f"Audio too long ({audio_len_sec:.1f}s). Please upload or record up to {MAX_AUDIO_SECONDS} seconds.")
        
        print("[LOG] Step 3: Extracting features...")
        inputs = processor(
            audio_array,
            sampling_rate=16000,
            return_tensors="pt"
        )
        del audio_array
        gc.collect()
        print(f"[LOG] Memory after preprocessing: {get_mem_usage():.2f} MB")
        
        print("[LOG] Step 4: Running model inference...")
        with torch.no_grad():
            outputs = model(input_features=inputs.input_features)
        
        del inputs
        gc.collect()
        
        print("[LOG] Step 5: Post-processing results...")
        fam_probs = torch.softmax(outputs["fam_logits"], dim=-1)
        super_probs = torch.softmax(outputs["super_logits"], dim=-1)
        code_probs = torch.softmax(outputs["code_logits"], dim=-1)
        
        def build_df(probs_tensor, k, thresh, labels_list, apply_mapping=False):
            k = int(k)
            top_vals, top_idx = torch.topk(probs_tensor[0], min(k, probs_tensor.shape[-1]))
            
            table_data = []
            for i in range(len(top_vals)):
                score = top_vals[i].item()
                if score < thresh:
                    continue
                
                idx = top_idx[i].item()
                raw_label = labels_list[idx].strip("<|>") if idx < len(labels_list) else f"Unknown ({idx})"
                
                if apply_mapping:
                    name = f"{CODE_TO_NAME[raw_label]} ({raw_label})" if raw_label in CODE_TO_NAME else raw_label
                else:
                    name = raw_label
                
                table_data.append([name, f"{score:.2%}"])
                
            if not table_data:
                return pd.DataFrame(columns=["Prediction / Predicción", "Confidence / Confianza"])
            return pd.DataFrame(table_data, columns=["Prediction / Predicción", "Confidence / Confianza"])

        df_fam = build_df(fam_probs, fam_k, fam_thresh, label_extractor.family_labels)
        df_super = build_df(super_probs, super_k, super_thresh, label_extractor.super_labels)
        df_code = build_df(code_probs, code_k, code_thresh, label_extractor.code_labels, apply_mapping=True)

        print(f"[LOG] Final Memory: {get_mem_usage():.2f} MB")
        print(f"--- [LOG] Request Finished ---\n")
        
        return df_fam, df_super, df_code
    except Exception as e:
        print(f"Error during inference: {e}")
        raise gr.Error(f"Processing failed: {str(e)}")

import base64
import re

# --- Load icon.png as base64 ---
try:
    with open("icon.png", "rb") as image_file:
        encoded_string = base64.b64encode(image_file.read()).decode("utf-8")
        icon_html = f'<img src="data:image/png;base64,{encoded_string}" style="height: 80px; position:absolute;" alt="Tachiwin Icon" />'
except Exception as e:
    icon_html = ""

# --- Load curated examples ---
curated_labels = []
curated_audio = []
try:
    with open("results.txt", "r", encoding="utf-8") as f:
        results_content = f.read()
    
    # Split by the sample header
    sample_blocks = re.split(r'📊 Sample (\d+)', results_content)
    for i in range(1, len(sample_blocks), 2):
        sample_num = sample_blocks[i]
        block = sample_blocks[i+1].strip()
        audio_path = f"samples/sample{sample_num}.wav"
        if os.path.exists(audio_path):
            curated_audio.append(audio_path)
            curated_labels.append(f"Sample {sample_num}:\n{block}")
except Exception as e:
    print(f"Warning: Could not parse results.txt: {e}")

# === UI COMPONENTS ===
with gr.Blocks(theme=gr.themes.Ocean()) as demo:
    gr.HTML(
        f"""
        <div style="text-align: center; padding: 30px; background: linear-gradient(120deg, rgb(2, 132, 199) 0%, rgb(16, 185, 129) 60%, rgb(5, 150, 105) 100%); color: white; border-radius: 15px; margin-bottom: 25px; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1);">
            {icon_html}
            <h1 style="color: white; margin: 0; font-size: 2.5em;">🦡 Tachiwin AudioId 🦡</h1>
            <p style="font-size: 1.2em; opacity: 0.9; margin-top: 10px;">Identify any of the 68 languages of Mexico and their 360 variants<br/>
            <i>Identifica cualquiera de las 68 lenguas originarias de México y sus 360 variantes</i></p>
        </div>
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(i18n("input_audio_header"))
            audio_input = gr.Audio(
                sources=["upload", "microphone"],
                type="filepath", # Changed from numpy to filepath
                label=i18n("upload_or_record_label")
            )
            with gr.Accordion(i18n("advanced_options"), open=False):
                with gr.Group():
                    gr.Markdown(i18n("language_family"))
                    with gr.Row():
                        fam_k = gr.Slider(minimum=1, maximum=10, step=1, value=1, label=i18n("top_k_label"))
                        fam_thresh = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.0, label=i18n("threshold_label"))
                with gr.Group():
                    gr.Markdown(i18n("superlanguage"))
                    with gr.Row():
                        super_k = gr.Slider(minimum=1, maximum=10, step=1, value=1, label=i18n("top_k_label"))
                        super_thresh = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.0, label=i18n("threshold_label"))
                with gr.Group():
                    gr.Markdown(i18n("language_code"))
                    with gr.Row():
                        code_k = gr.Slider(minimum=1, maximum=10, step=1, value=3, label=i18n("top_k_label"))
                        code_thresh = gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.0, label=i18n("threshold_label"))
                        
            with gr.Row():
                clear_btn = gr.Button(i18n("clear_btn"), variant="secondary")
                submit_btn = gr.Button(i18n("classify_btn"), variant="primary")
            
        with gr.Column(scale=1):
            gr.Markdown(i18n("results_header"))
            fam_table = gr.Dataframe(headers=["Prediction / Predicción", "Confidence / Confianza"], datatype=["str", "str"], label=i18n("pred_fam_label"), interactive=False, wrap=True)
            super_table = gr.Dataframe(headers=["Prediction / Predicción", "Confidence / Confianza"], datatype=["str", "str"], label=i18n("pred_super_label"), interactive=False, wrap=True)
            code_table = gr.Dataframe(headers=["Prediction / Predicción", "Confidence / Confianza"], datatype=["str", "str"], label=i18n("pred_code_label"), interactive=False, wrap=True)

    submit_btn.click(
        fn=predict_language,
        inputs=[audio_input, fam_k, fam_thresh, super_k, super_thresh, code_k, code_thresh],
        outputs=[fam_table, super_table, code_table]
    )
    
    clear_btn.click(
        fn=lambda: (None, None, None, None),
        inputs=None,
        outputs=[audio_input, fam_table, super_table, code_table]
    )

    gr.Markdown(i18n("curated_header"))
    gr.Markdown(i18n("curated_desc"))
    
    with gr.Group():
        columns = 3
        for row_idx in range(0, len(curated_audio), columns):
            with gr.Row():
                for i in range(row_idx, min(row_idx+columns, len(curated_audio))):
                    with gr.Column(scale=1):
                        gr.Markdown(f"```text\n{curated_labels[i]}\n```")
                        gr.Audio(value=curated_audio[i], interactive=False, show_label=False)
                        btn = gr.Button("Load / Cargar", size="sm")
                        btn.click(fn=lambda url=curated_audio[i]: url, inputs=[], outputs=[audio_input])

if __name__ == "__main__":
    demo.launch(ssr_mode=False, i18n=i18n)