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'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"""
{icon_html}

🦡 Tachiwin AudioId 🦡

Identify any of the 68 languages of Mexico and their 360 variants
Identifica cualquiera de las 68 lenguas originarias de México y sus 360 variantes

""" ) 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)