#!/usr/bin/env python3
"""
Voxtral ASR Fine-tuning Interface
Features:
- Collect a personal voice dataset (upload WAV/FLAC + transcripts or record mic audio)
- Build a JSONL dataset ({audio_path, text}) at 16kHz
- Fine-tune Voxtral (LoRA or full) with streamed logs
- Push model to Hugging Face Hub
- Deploy a Voxtral ASR demo Space
Env tokens (optional):
- HF_WRITE_TOKEN or HF_TOKEN: write access token
- HF_READ_TOKEN: optional read token
- HF_USERNAME: fallback username if not derivable from token
"""
from __future__ import annotations
import os
import json
from pathlib import Path
from datetime import datetime
from typing import Any, Dict, Generator, Optional, Tuple
import gradio as gr
PROJECT_ROOT = Path(__file__).resolve().parent
def get_python() -> str:
import sys
return sys.executable or "python"
def get_username_from_token(token: str) -> Optional[str]:
try:
from huggingface_hub import HfApi # type: ignore
api = HfApi(token=token)
info = api.whoami()
if isinstance(info, dict):
return info.get("name") or info.get("username")
if isinstance(info, str):
return info
except Exception:
return None
return None
def run_command_stream(args: list[str], env: Dict[str, str], cwd: Optional[Path] = None) -> Generator[str, None, int]:
import subprocess
import shlex
try:
cmd_line = ' '.join(shlex.quote(a) for a in ([get_python()] + args))
yield f"$ {cmd_line}"
process = subprocess.Popen(
[get_python()] + args,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
env=env,
cwd=str(cwd or PROJECT_ROOT),
bufsize=1,
universal_newlines=True,
)
if process.stdout is None:
yield "❌ Error: Could not capture process output"
return 1
for line in iter(process.stdout.readline, ""):
if line.strip(): # Only yield non-empty lines
yield line.rstrip()
process.stdout.close()
code = process.wait()
if code != 0:
yield f"❌ Command failed with exit code: {code}"
else:
yield f"✅ Command completed successfully (exit code: {code})"
return code
except FileNotFoundError as e:
yield f"❌ Error: Python executable not found: {e}"
return 1
except Exception as e:
yield f"❌ Error running command: {str(e)}"
return 1
def detect_nvidia_driver() -> Tuple[bool, str]:
"""Detect NVIDIA driver/GPU presence with multiple strategies.
Returns (available, human_message).
"""
# 1) Try torch CUDA
try:
import torch # type: ignore
if torch.cuda.is_available():
try:
num = torch.cuda.device_count()
names = [torch.cuda.get_device_name(i) for i in range(num)]
return True, f"NVIDIA GPU detected: {', '.join(names)}"
except Exception:
return True, "NVIDIA GPU detected (torch.cuda available)"
except Exception:
pass
# 2) Try NVML via pynvml
try:
import pynvml # type: ignore
try:
pynvml.nvmlInit()
cnt = pynvml.nvmlDeviceGetCount()
names = []
for i in range(cnt):
h = pynvml.nvmlDeviceGetHandleByIndex(i)
names.append(pynvml.nvmlDeviceGetName(h).decode("utf-8", errors="ignore"))
drv = pynvml.nvmlSystemGetDriverVersion().decode("utf-8", errors="ignore")
pynvml.nvmlShutdown()
if cnt > 0:
return True, f"NVIDIA driver {drv}; GPUs: {', '.join(names)}"
except Exception:
pass
except Exception:
pass
# 3) Try nvidia-smi
try:
import subprocess
res = subprocess.run(["nvidia-smi", "-L"], capture_output=True, text=True, timeout=3)
if res.returncode == 0 and res.stdout.strip():
return True, res.stdout.strip().splitlines()[0]
except Exception:
pass
return False, "No NVIDIA driver/GPU detected"
def duplicate_space_hint() -> str:
space_id = os.environ.get("SPACE_ID") or os.environ.get("HF_SPACE_ID")
if space_id:
space_url = f"https://huggingface.co/spaces/{space_id}"
dup_url = f"{space_url}?duplicate=true"
return (
f"ℹ️ No NVIDIA driver detected. If you're on Hugging Face Spaces, "
f"please duplicate this Space to GPU hardware: [Duplicate this Space]({dup_url})."
)
return (
"ℹ️ No NVIDIA driver detected. To enable training, run on a machine with an NVIDIA GPU/driver "
"or duplicate this Space on Hugging Face with GPU hardware."
)
def _write_jsonl(rows: list[dict], path: Path) -> Path:
path.parent.mkdir(parents=True, exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
for r in rows:
f.write(json.dumps(r, ensure_ascii=False) + "\n")
return path
def _save_uploaded_dataset(files: list, transcripts: list[str]) -> str:
dataset_dir = PROJECT_ROOT / "datasets" / "voxtral_user"
dataset_dir.mkdir(parents=True, exist_ok=True)
rows: list[dict] = []
for i, fpath in enumerate(files or []):
if i >= len(transcripts):
break
rows.append({"audio_path": fpath, "text": transcripts[i] or ""})
jsonl_path = dataset_dir / "data.jsonl"
_write_jsonl(rows, jsonl_path)
return str(jsonl_path)
def _push_dataset_to_hub(jsonl_path: str, repo_name: str, username: str = "") -> str:
"""Push dataset to Hugging Face Hub including audio files"""
try:
from huggingface_hub import HfApi, create_repo
import json
from pathlib import Path
import os
token = os.getenv("HF_TOKEN") or os.getenv("HF_WRITE_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
if not token:
return "❌ No HF_TOKEN found. Set HF_TOKEN environment variable to push datasets."
api = HfApi(token=token)
# Determine full repo name
if "/" not in repo_name:
if not username:
user_info = api.whoami()
username = user_info.get("name") or user_info.get("username") or ""
if username:
repo_name = f"{username}/{repo_name}"
# Create dataset repository
try:
create_repo(repo_name, repo_type="dataset", token=token, exist_ok=True)
except Exception as e:
if "already exists" not in str(e).lower():
return f"❌ Failed to create dataset repo: {e}"
# Read the JSONL file
jsonl_file = Path(jsonl_path)
if not jsonl_file.exists():
return f"❌ Dataset file not found: {jsonl_path}"
# Read and process the JSONL to collect audio files and update paths
audio_files = []
updated_rows = []
total_audio_size = 0
with open(jsonl_file, "r", encoding="utf-8") as f:
for line_num, line in enumerate(f):
try:
row = json.loads(line.strip())
audio_path = row.get("audio_path", "")
if audio_path:
audio_file = Path(audio_path)
if audio_file.exists():
# Store the original file for upload
audio_files.append(audio_file)
total_audio_size += audio_file.stat().st_size
# Update path to be relative for the dataset
row["audio_path"] = f"audio/{audio_file.name}"
else:
print(f"⚠️ Warning: Audio file not found: {audio_path}")
row["audio_path"] = "" # Clear missing files
updated_rows.append(row)
except json.JSONDecodeError as e:
print(f"⚠️ Warning: Invalid JSON on line {line_num + 1}: {e}")
continue
# Create updated JSONL with relative paths
temp_jsonl_path = jsonl_file.parent / "temp_data.jsonl"
with open(temp_jsonl_path, "w", encoding="utf-8") as f:
for row in updated_rows:
f.write(json.dumps(row, ensure_ascii=False) + "\n")
# Upload the updated JSONL file
api.upload_file(
path_or_fileobj=str(temp_jsonl_path),
path_in_repo="data.jsonl",
repo_id=repo_name,
repo_type="dataset",
token=token
)
# Clean up temp file
temp_jsonl_path.unlink()
# Upload audio files
uploaded_count = 0
for audio_file in audio_files:
try:
remote_path = f"audio/{audio_file.name}"
api.upload_file(
path_or_fileobj=str(audio_file),
path_in_repo=remote_path,
repo_id=repo_name,
repo_type="dataset",
token=token
)
uploaded_count += 1
print(f"✅ Uploaded audio file: {audio_file.name}")
except Exception as e:
print(f"❌ Failed to upload {audio_file.name}: {e}")
# Calculate total dataset size
total_dataset_size = jsonl_file.stat().st_size + total_audio_size
# Create README for the dataset
readme_content = f"""---
dataset_info:
features:
- name: audio_path
dtype: string
- name: text
dtype: string
splits:
- name: train
num_bytes: {jsonl_file.stat().st_size}
num_examples: {len(updated_rows)}
download_size: {total_dataset_size}
dataset_size: {total_dataset_size}
tags:
- voxtral
- asr
- speech-to-text
- fine-tuning
- audio-dataset
---
# Voxtral ASR Dataset
This dataset was created using the Voxtral ASR Fine-tuning Interface.
## Dataset Structure
- **audio_path**: Relative path to the audio file (stored in `audio/` directory)
- **text**: Transcription of the audio
## Dataset Statistics
- **Number of examples**: {len(updated_rows)}
- **Audio files uploaded**: {uploaded_count}
- **Total dataset size**: {total_dataset_size:,} bytes
## Usage
```python
from datasets import load_dataset, Audio
# Load dataset
dataset = load_dataset("{repo_name}")
# Load audio data
dataset = dataset.cast_column("audio_path", Audio())
# Access first example
print(dataset[0]["text"])
print(dataset[0]["audio_path"])
```
## Loading with Audio Decoding
```python
from datasets import load_dataset, Audio
# Load with automatic audio decoding
dataset = load_dataset("{repo_name}")
dataset = dataset.cast_column("audio_path", Audio(sampling_rate=16000))
# The audio column will contain the decoded audio arrays
audio_array = dataset[0]["audio_path"]["array"]
sampling_rate = dataset[0]["audio_path"]["sampling_rate"]
```
"""
# Upload README
readme_path = jsonl_file.parent / "README.md"
with open(readme_path, "w") as f:
f.write(readme_content)
api.upload_file(
path_or_fileobj=str(readme_path),
path_in_repo="README.md",
repo_id=repo_name,
repo_type="dataset",
token=token
)
readme_path.unlink() # Clean up temp file
return f"✅ Dataset pushed to: https://huggingface.co/datasets/{repo_name}\n📊 Uploaded {len(updated_rows)} examples and {uploaded_count} audio files"
except Exception as e:
return f"❌ Failed to push dataset: {e}"
def _save_recordings(recordings: list[tuple[int, list]], transcripts: list[str]) -> str:
import soundfile as sf
dataset_dir = PROJECT_ROOT / "datasets" / "voxtral_user"
wav_dir = dataset_dir / "wavs"
wav_dir.mkdir(parents=True, exist_ok=True)
rows: list[dict] = []
for i, rec in enumerate(recordings or []):
if rec is None:
continue
if i >= len(transcripts):
break
sr, data = rec
out_path = wav_dir / f"rec_{i:04d}.wav"
sf.write(str(out_path), data, sr)
rows.append({"audio_path": str(out_path), "text": transcripts[i] or ""})
jsonl_path = dataset_dir / "data.jsonl"
_write_jsonl(rows, jsonl_path)
return str(jsonl_path)
def start_voxtral_training(
use_lora: bool,
base_model: str,
repo_short: str,
jsonl_path: str,
train_count: int,
eval_count: int,
batch_size: int,
grad_accum: int,
learning_rate: float,
epochs: float,
lora_r: int,
lora_alpha: int,
lora_dropout: float,
freeze_audio_tower: bool,
push_to_hub: bool,
deploy_demo: bool,
) -> str:
"""Start Voxtral training and return collected logs as a string."""
env = os.environ.copy()
write_token = env.get("HF_WRITE_TOKEN") or env.get("HF_TOKEN")
read_token = env.get("HF_READ_TOKEN")
username = get_username_from_token(write_token or "") or env.get("HF_USERNAME") or ""
output_dir = PROJECT_ROOT / "outputs" / repo_short
# Collect all logs
all_logs = []
def collect_logs(generator):
"""Helper to collect logs from a generator."""
for line in generator:
all_logs.append(line)
print(line) # Also print to console for debugging
try:
# 1) Train
script = PROJECT_ROOT / ("scripts/train_lora.py" if use_lora else "scripts/train.py")
args = [str(script)]
if jsonl_path:
args += ["--dataset-jsonl", jsonl_path]
args += [
"--model-checkpoint", base_model,
"--train-count", str(train_count),
"--eval-count", str(eval_count),
"--batch-size", str(batch_size),
"--grad-accum", str(grad_accum),
"--learning-rate", str(learning_rate),
"--epochs", str(epochs),
"--output-dir", str(output_dir),
"--save-steps", "50",
]
if use_lora:
args += [
"--lora-r", str(lora_r),
"--lora-alpha", str(lora_alpha),
"--lora-dropout", str(lora_dropout),
]
if freeze_audio_tower:
args += ["--freeze-audio-tower"]
all_logs.append("🚀 Starting Voxtral training...")
collect_logs(run_command_stream(args, env))
all_logs.append("✅ Training completed!")
# 2) Push to Hub
if push_to_hub:
if not username:
all_logs.append("❌ Cannot push to Hub: No username available. Set HF_TOKEN or HF_USERNAME.")
else:
repo_name = f"{username}/{repo_short}"
push_args = [
str(PROJECT_ROOT / "scripts/push_to_huggingface.py"),
"model",
str(output_dir),
repo_name,
]
all_logs.append(f"📤 Pushing model to Hugging Face Hub: {repo_name}")
collect_logs(run_command_stream(push_args, env))
all_logs.append("✅ Model pushed successfully!")
# 3) Deploy demo Space
if deploy_demo and username:
deploy_args = [
str(PROJECT_ROOT / "scripts/deploy_demo_space.py"),
"--hf-token", write_token or "",
"--hf-username", username,
"--model-id", f"{username}/{repo_short}",
"--demo-type", "voxtral",
"--space-name", f"{repo_short}-demo",
]
all_logs.append("🚀 Deploying demo Space...")
collect_logs(run_command_stream(deploy_args, env))
all_logs.append("✅ Demo Space deployed!")
# Return all collected logs as a single string
return "\n".join(all_logs)
except Exception as e:
error_msg = f"❌ Error during training: {str(e)}"
all_logs.append(error_msg)
print(error_msg) # Also print to console
import traceback
traceback.print_exc()
return "\n".join(all_logs)
def load_multilingual_phrases(language="en", max_phrases=None, split="train"):
"""Load phrases from NVIDIA Granary dataset.
Uses the high-quality Granary dataset which contains speech recognition
and translation data for 25 European languages.
Args:
language: Language code (e.g., 'en', 'de', 'fr', etc.)
max_phrases: Maximum number of phrases to load (None for default 1000)
split: Dataset split to use ('train', 'validation', 'test')
Returns:
List of transcription phrases from Granary dataset
"""
from datasets import load_dataset
import random
# Default to 1000 phrases if not specified
if max_phrases is None:
max_phrases = 1000
# Language code mapping for CohereLabs AYA Collection dataset
# All Voxtral Mini supported languages are available in AYA Collection
aya_supported_langs = {
"en": "english", # English
"fr": "french", # French
"de": "german", # German
"es": "spanish", # Spanish
"it": "italian", # Italian
"pt": "portuguese", # Portuguese
"nl": "dutch", # Dutch
"hi": "hindi" # Hindi
}
# Map input language to CohereLabs AYA Collection configuration
aya_lang = aya_supported_langs.get(language)
if not aya_lang:
raise Exception(f"Language {language} not supported in CohereLabs AYA Collection dataset")
try:
print(f"Loading phrases from CohereLabs AYA Collection dataset for language: {language}")
# Check for authentication token
token = os.getenv("HF_TOKEN") or os.getenv("HF_WRITE_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
# Try to load CohereLabs AYA Collection dataset for the specified language
if token:
try:
ds = load_dataset("CohereLabs/aya_collection_language_split", aya_lang, split="train", streaming=True, token=token)
print(f"Successfully loaded CohereLabs AYA Collection {language} dataset")
except Exception as e:
# Fallback to other datasets
print(f"CohereLabs AYA Collection {language} not available ({e}), trying alternative datasets")
raise Exception("AYA Collection not available")
else:
print("No HF_TOKEN found for CohereLabs AYA Collection dataset")
raise Exception("No token available")
# Common processing for both dataset types
phrases = []
count = 0
seen_phrases = set()
# Sample phrases from the dataset
for example in ds:
if count >= max_phrases:
break
# Extract text from CohereLabs AYA Collection format: combine inputs and targets
inputs_text = example.get("inputs", "").strip()
targets_text = example.get("targets", "").strip()
text = f"{inputs_text} {targets_text}".strip()
# Filter for quality phrases
if (text and
len(text) > 10 and # Minimum length
len(text) < 200 and # Maximum length to avoid very long utterances
text not in seen_phrases and # Avoid duplicates
not text.isdigit() and # Avoid pure numbers
not all(c in "0123456789., " for c in text)): # Avoid mostly numeric
phrases.append(text)
seen_phrases.add(text)
count += 1
if phrases:
# Shuffle the phrases for variety
random.shuffle(phrases)
dataset_name = "CohereLabs AYA Collection"
print(f"Successfully loaded {len(phrases)} phrases from {dataset_name} dataset for {language}")
return phrases
else:
print(f"No suitable phrases found in dataset for {language}")
raise Exception("No phrases found")
except Exception as e:
error_msg = str(e).lower()
if "401" in error_msg or "unauthorized" in error_msg:
print(f"CohereLabs AYA Collection authentication failed for {language}: {e}")
print("This dataset requires a Hugging Face token. Please set HF_TOKEN environment variable.")
else:
print(f"CohereLabs AYA Collection loading failed for {language}: {e}")
# Fallback to basic phrases if dataset loading fails
print("Using fallback phrases")
# Language-specific fallback phrases
language_fallbacks = {
"hi": [
"नमस्ते, आज आप कैसे हैं?",
"मेरा नाम राजेश कुमार है।",
"आज का मौसम बहुत अच्छा है।",
"मैं हिंदी में बात करना चाहता हूं।",
"कृपया धीरे और स्पष्ट बोलें।",
"यह एक परीक्षण वाक्य है।",
"मैं पुस्तकें पढ़ना पसंद करता हूं।",
"क्या आप मेरी मदद कर सकते हैं?",
"आपका फोन नंबर क्या है?",
"मैं कल सुबह आऊंगा।",
"धन्यवाद, आपका समय देने के लिए।",
"यह जगह बहुत सुंदर है।",
"मैं भोजन तैयार करना सीख रहा हूं।",
"क्या यह रास्ता सही है?",
"मैं स्कूल जाना चाहता हूं।",
"आपकी उम्र क्या है?",
"यह कितने का है?",
"मैं थक गया हूं।",
"आप कहां से हैं?",
"चलिए पार्क में टहलते हैं।"
],
"en": [
"Hello, how are you today?",
"My name is John Smith.",
"The weather is very nice today.",
"I want to speak in English.",
"Please speak slowly and clearly.",
"This is a test sentence.",
"I enjoy reading books.",
"Can you help me?",
"What is your phone number?",
"I will come tomorrow morning.",
"Thank you for your time.",
"This place is very beautiful.",
"I am learning to cook food.",
"Is this the right way?",
"I want to go to school.",
"How old are you?",
"How much does this cost?",
"I am tired.",
"Where are you from?",
"Let's go for a walk in the park."
],
"fr": [
"Bonjour, comment allez-vous aujourd'hui?",
"Je m'appelle Jean Dupont.",
"Le temps est très beau aujourd'hui.",
"Je veux parler en français.",
"Parlez lentement et clairement s'il vous plaît.",
"Ceci est une phrase de test.",
"J'aime lire des livres.",
"Pouvez-vous m'aider?",
"Quel est votre numéro de téléphone?",
"Je viendrai demain matin.",
"Merci pour votre temps.",
"Cet endroit est très beau.",
"J'apprends à cuisiner.",
"Est-ce le bon chemin?",
"Je veux aller à l'école.",
"Quel âge avez-vous?",
"Combien cela coûte-t-il?",
"Je suis fatigué.",
"D'où venez-vous?",
"Allons nous promener dans le parc."
],
"de": [
"Hallo, wie geht es Ihnen heute?",
"Mein Name ist Hans Müller.",
"Das Wetter ist heute sehr schön.",
"Ich möchte auf Deutsch sprechen.",
"Sprechen Sie bitte langsam und deutlich.",
"Dies ist ein Testsatz.",
"Ich lese gerne Bücher.",
"Können Sie mir helfen?",
"Wie ist Ihre Telefonnummer?",
"Ich komme morgen früh.",
"Vielen Dank für Ihre Zeit.",
"Dieser Ort ist sehr schön.",
"Ich lerne kochen.",
"Ist das der richtige Weg?",
"Ich möchte zur Schule gehen.",
"Wie alt sind Sie?",
"Wie viel kostet das?",
"Ich bin müde.",
"Woher kommen Sie?",
"Lassen Sie uns im Park spazieren gehen."
],
"es": [
"Hola, ¿cómo estás hoy?",
"Me llamo Juan García.",
"El tiempo está muy bueno hoy.",
"Quiero hablar en español.",
"Por favor habla despacio y claro.",
"Esta es una oración de prueba.",
"Me gusta leer libros.",
"¿Puedes ayudarme?",
"¿Cuál es tu número de teléfono?",
"Vendré mañana por la mañana.",
"Gracias por tu tiempo.",
"Este lugar es muy bonito.",
"Estoy aprendiendo a cocinar.",
"¿Es este el camino correcto?",
"Quiero ir a la escuela.",
"¿Cuántos años tienes?",
"¿Cuánto cuesta esto?",
"Estoy cansado.",
"¿De dónde eres?",
"Vamos a caminar por el parque."
],
"it": [
"Ciao, come stai oggi?",
"Mi chiamo Mario Rossi.",
"Il tempo è molto bello oggi.",
"Voglio parlare in italiano.",
"Per favore parla lentamente e chiaramente.",
"Questa è una frase di prova.",
"Mi piace leggere libri.",
"Puoi aiutarmi?",
"Qual è il tuo numero di telefono?",
"Verrò domani mattina.",
"Grazie per il tuo tempo.",
"Questo posto è molto bello.",
"Sto imparando a cucinare.",
"È questa la strada giusta?",
"Voglio andare a scuola.",
"Quanti anni hai?",
"Quanto costa questo?",
"Sono stanco.",
"Da dove vieni?",
"Andiamo a fare una passeggiata nel parco."
],
"pt": [
"Olá, como você está hoje?",
"Meu nome é João Silva.",
"O tempo está muito bom hoje.",
"Quero falar em português.",
"Por favor fale devagar e claramente.",
"Esta é uma frase de teste.",
"Eu gosto de ler livros.",
"Você pode me ajudar?",
"Qual é o seu número de telefone?",
"Vou vir amanhã de manhã.",
"Obrigado pelo seu tempo.",
"Este lugar é muito bonito.",
"Estou aprendendo a cozinhar.",
"Este é o caminho certo?",
"Quero ir para a escola.",
"Quantos anos você tem?",
"Quanto custa isso?",
"Estou cansado.",
"De onde você é?",
"Vamos dar um passeio no parque."
],
"nl": [
"Hallo, hoe gaat het vandaag met je?",
"Mijn naam is Jan de Vries.",
"Het weer is vandaag erg mooi.",
"Ik wil in het Nederlands spreken.",
"Spreek langzaam en duidelijk alstublieft.",
"Dit is een testzin.",
"Ik houd van het lezen van boeken.",
"Kun je me helpen?",
"Wat is je telefoonnummer?",
"Ik kom morgenochtend.",
"Bedankt voor je tijd.",
"Deze plek is erg mooi.",
"Ik leer koken.",
"Is dit de juiste weg?",
"Ik wil naar school gaan.",
"Hoe oud ben je?",
"Hoeveel kost dit?",
"Ik ben moe.",
"Waar kom je vandaan?",
"Laten we een wandeling maken in het park."
]
}
fallback_phrases = language_fallbacks.get(language, language_fallbacks["en"])
if max_phrases:
fallback_phrases = random.sample(fallback_phrases, min(max_phrases, len(fallback_phrases)))
else:
random.shuffle(fallback_phrases)
return fallback_phrases
# Initialize phrases dynamically
DEFAULT_LANGUAGE = "en" # Default to English
ALL_PHRASES = load_multilingual_phrases(DEFAULT_LANGUAGE, max_phrases=None)
with gr.Blocks(title="Voxtral ASR Fine-tuning") as demo:
has_gpu, gpu_msg = detect_nvidia_driver()
if has_gpu:
gr.HTML(
f"""
✅ NVIDIA GPU ready — {gpu_msg}
Set HF_WRITE_TOKEN/HF_TOKEN in environment to enable Hub push.
"""
)
else:
hint_md = duplicate_space_hint()
gr.HTML(
f"""
⚠️ No NVIDIA GPU/driver detected — training requires a GPU runtime
{hint_md}
"""
)
gr.Markdown("""
# 🎙️ Voxtral ASR Fine-tuning
Read the phrases below and record them. Then start fine-tuning.
""")
# Check for HF_TOKEN and show warning if missing
hf_token = os.getenv("HF_TOKEN") or os.getenv("HF_WRITE_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")
if not hf_token:
gr.HTML(
"""
⚠️ No HF_TOKEN detected
Set HF_TOKEN environment variable to access CohereLabs AYA Collection dataset with authentic multilingual phrases.
This dataset provides high-quality text in 100+ languages for all Voxtral Mini supported languages.
Currently using fallback phrases for demonstration.
"""
)
# Hidden state to track dataset JSONL path
jsonl_path_state = gr.State("")
# Language selection for Voxtral Mini supported languages
language_selector = gr.Dropdown(
choices=[
("English", "en"),
("French", "fr"),
("German", "de"),
("Spanish", "es"),
("Italian", "it"),
("Portuguese", "pt"),
("Dutch", "nl"),
("Hindi", "hi")
],
value="en",
label="Language for Speech Phrases",
info="Select language for authentic phrases (Voxtral Mini supported languages). All languages use CohereLabs AYA Collection dataset when HF_TOKEN is available."
)
# Recording grid with dynamic text readouts
phrase_texts_state = gr.State(ALL_PHRASES)
visible_rows_state = gr.State(10) # Start with 10 visible rows
MAX_COMPONENTS = 100 # Fixed maximum number of components
# Create fixed number of components upfront
phrase_markdowns: list[gr.Markdown] = []
rec_components = []
def create_recording_grid(max_components=MAX_COMPONENTS):
"""Create recording grid components with fixed maximum"""
markdowns = []
recordings = []
for idx in range(max_components):
visible = False # Initially hidden - will be revealed when language is selected
phrase_text = ALL_PHRASES[idx] if idx < len(ALL_PHRASES) else ""
md = gr.Markdown(f"**{idx+1}. {phrase_text}**", visible=visible)
markdowns.append(md)
comp = gr.Audio(sources="microphone", type="numpy", label=f"Recording {idx+1}", visible=visible)
recordings.append(comp)
return markdowns, recordings
# Initial grid creation
with gr.Column():
phrase_markdowns, rec_components = create_recording_grid(MAX_COMPONENTS)
# Add more rows button
add_rows_btn = gr.Button("➕ Add 10 More Rows", variant="secondary", visible=False)
def add_more_rows(current_visible, current_phrases):
"""Add 10 more rows by making them visible"""
new_visible = min(current_visible + 10, MAX_COMPONENTS, len(current_phrases))
# Create updates for all MAX_COMPONENTS (both markdown and audio components)
markdown_updates = []
audio_updates = []
for i in range(MAX_COMPONENTS):
if i < len(current_phrases) and i < new_visible:
markdown_updates.append(gr.update(visible=True))
audio_updates.append(gr.update(visible=True))
else:
markdown_updates.append(gr.update(visible=False))
audio_updates.append(gr.update(visible=False))
# Return: [state] + markdown_updates + audio_updates
return [new_visible] + markdown_updates + audio_updates
def change_language(language):
"""Change the language and reload phrases from multilingual datasets, reveal interface"""
new_phrases = load_multilingual_phrases(language, max_phrases=None)
# Reset visible rows to 10
visible_count = min(10, len(new_phrases), MAX_COMPONENTS)
# Create separate updates for markdown and audio components
markdown_updates = []
audio_updates = []
for i in range(MAX_COMPONENTS):
if i < len(new_phrases) and i < visible_count:
markdown_updates.append(gr.update(value=f"**{i+1}. {new_phrases[i]}**", visible=True))
audio_updates.append(gr.update(visible=True))
elif i < len(new_phrases):
markdown_updates.append(gr.update(value=f"**{i+1}. {new_phrases[i]}**", visible=False))
audio_updates.append(gr.update(visible=False))
else:
markdown_updates.append(gr.update(value=f"**{i+1}. **", visible=False))
audio_updates.append(gr.update(visible=False))
# Reveal all interface elements when language is selected
reveal_updates = [
gr.update(visible=True), # add_rows_btn
gr.update(visible=True), # record_dataset_btn
gr.update(visible=True), # dataset_status
gr.update(visible=True), # advanced_accordion
gr.update(visible=True), # save_rec_btn
gr.update(visible=True), # push_recordings_btn
gr.update(visible=True), # start_btn
gr.update(visible=True), # logs_box
]
# Return: [phrases_state, visible_state] + markdown_updates + audio_updates + reveal_updates
return [new_phrases, visible_count] + markdown_updates + audio_updates + reveal_updates
add_rows_btn.click(
add_more_rows,
inputs=[visible_rows_state, phrase_texts_state],
outputs=[visible_rows_state] + phrase_markdowns + rec_components
)
# Recording dataset creation button
record_dataset_btn = gr.Button("🎙️ Create Dataset from Recordings", variant="primary", visible=False)
def create_recording_dataset(*recordings_and_state):
"""Create dataset from visible recordings and phrases"""
try:
import soundfile as sf
# Extract recordings and state
recordings = recordings_and_state[:-1] # All except the last item (phrases)
phrases = recordings_and_state[-1] # Last item is phrases
dataset_dir = PROJECT_ROOT / "datasets" / "voxtral_user"
wav_dir = dataset_dir / "wavs"
wav_dir.mkdir(parents=True, exist_ok=True)
rows = []
successful_recordings = 0
# Process each recording
for i, rec in enumerate(recordings):
if rec is not None and i < len(phrases):
try:
sr, data = rec
out_path = wav_dir / f"recording_{i:04d}.wav"
sf.write(str(out_path), data, sr)
rows.append({"audio_path": str(out_path), "text": phrases[i]})
successful_recordings += 1
except Exception as e:
print(f"Error processing recording {i}: {e}")
if rows:
jsonl_path = dataset_dir / "recorded_data.jsonl"
_write_jsonl(rows, jsonl_path)
return f"✅ Dataset created successfully! {successful_recordings} recordings saved to {jsonl_path}"
else:
return "❌ No recordings found. Please record some audio first."
except Exception as e:
return f"❌ Error creating dataset: {str(e)}"
# Status display for dataset creation
dataset_status = gr.Textbox(label="Dataset Creation Status", interactive=False, visible=False)
record_dataset_btn.click(
create_recording_dataset,
inputs=rec_components + [phrase_texts_state],
outputs=[dataset_status]
)
# Advanced options accordion
with gr.Accordion("Advanced options", open=False, visible=False) as advanced_accordion:
base_model = gr.Textbox(value="mistralai/Voxtral-Mini-3B-2507", label="Base Voxtral model")
use_lora = gr.Checkbox(value=True, label="Use LoRA (parameter-efficient)")
with gr.Row():
batch_size = gr.Number(value=2, precision=0, label="Batch size")
grad_accum = gr.Number(value=4, precision=0, label="Grad accum")
with gr.Row():
learning_rate = gr.Number(value=5e-5, precision=6, label="Learning rate")
epochs = gr.Number(value=3.0, precision=2, label="Epochs")
with gr.Accordion("LoRA settings", open=False):
lora_r = gr.Number(value=8, precision=0, label="LoRA r")
lora_alpha = gr.Number(value=32, precision=0, label="LoRA alpha")
lora_dropout = gr.Number(value=0.0, precision=3, label="LoRA dropout")
freeze_audio_tower = gr.Checkbox(value=True, label="Freeze audio tower")
with gr.Row():
train_count = gr.Number(value=100, precision=0, label="Train samples")
eval_count = gr.Number(value=50, precision=0, label="Eval samples")
repo_short = gr.Textbox(value=f"voxtral-finetune-{datetime.now().strftime('%Y%m%d_%H%M%S')}", label="Model repo (short)")
push_to_hub = gr.Checkbox(value=True, label="Push to HF Hub after training")
deploy_demo = gr.Checkbox(value=True, label="Deploy demo Space after push")
gr.Markdown("### Upload audio + transcripts (optional)")
upload_audio = gr.File(file_count="multiple", type="filepath", label="Upload WAV/FLAC files (optional)")
transcripts_box = gr.Textbox(lines=6, label="Transcripts (one per line, aligned with files)")
dataset_repo_name = gr.Textbox(value=f"voxtral-dataset-{datetime.now().strftime('%Y%m%d_%H%M%S')}",
label="Dataset repo name (will be pushed to HF Hub)")
save_upload_btn = gr.Button("Save uploaded dataset")
push_dataset_btn = gr.Button("Push dataset to HF Hub")
def _collect_upload(files, txt):
lines = [s.strip() for s in (txt or "").splitlines() if s.strip()]
jsonl_path = _save_uploaded_dataset(files or [], lines)
return str(jsonl_path), f"✅ Dataset saved locally: {jsonl_path}"
def _push_dataset_handler(repo_name, current_jsonl_path):
if not current_jsonl_path:
return "❌ No dataset saved yet. Please save dataset first."
return _push_dataset_to_hub(current_jsonl_path, repo_name)
save_upload_btn.click(_collect_upload, [upload_audio, transcripts_box], [jsonl_path_state, dataset_status])
push_dataset_btn.click(_push_dataset_handler, [dataset_repo_name, jsonl_path_state], [dataset_status])
# Save recordings button
save_rec_btn = gr.Button("Save recordings as dataset", visible=False)
push_recordings_btn = gr.Button("Push recordings dataset to HF Hub", visible=False)
def _collect_preloaded_recs(*recs_and_texts):
import soundfile as sf
dataset_dir = PROJECT_ROOT / "datasets" / "voxtral_user"
wav_dir = dataset_dir / "wavs"
wav_dir.mkdir(parents=True, exist_ok=True)
rows: list[dict] = []
if not recs_and_texts:
jsonl_path = dataset_dir / "data.jsonl"
_write_jsonl(rows, jsonl_path)
return str(jsonl_path)
texts = recs_and_texts[-1]
recs = recs_and_texts[:-1]
for i, rec in enumerate(recs):
if rec is None:
continue
sr, data = rec
out_path = wav_dir / f"rec_{i:04d}.wav"
sf.write(str(out_path), data, sr)
# Use the full phrase list (ALL_PHRASES) instead of just PHRASES
label_text = (texts[i] if isinstance(texts, list) and i < len(texts) else (ALL_PHRASES[i] if i < len(ALL_PHRASES) else ""))
rows.append({"audio_path": str(out_path), "text": label_text})
jsonl_path = dataset_dir / "data.jsonl"
_write_jsonl(rows, jsonl_path)
return str(jsonl_path), f"✅ Dataset saved locally: {jsonl_path}"
save_rec_btn.click(_collect_preloaded_recs, rec_components + [phrase_texts_state], [jsonl_path_state, dataset_status])
def _push_recordings_handler(repo_name, current_jsonl_path):
if not current_jsonl_path:
return "❌ No recordings dataset saved yet. Please save recordings first."
return _push_dataset_to_hub(current_jsonl_path, repo_name)
push_recordings_btn.click(_push_recordings_handler, [dataset_repo_name, jsonl_path_state], [dataset_status])
# Removed multilingual dataset sample section - phrases are now loaded automatically when language is selected
start_btn = gr.Button("Start Fine-tuning", visible=False)
logs_box = gr.Textbox(label="Logs", lines=20, visible=False)
start_btn.click(
start_voxtral_training,
inputs=[
use_lora, base_model, repo_short, jsonl_path_state, train_count, eval_count,
batch_size, grad_accum, learning_rate, epochs,
lora_r, lora_alpha, lora_dropout, freeze_audio_tower,
push_to_hub, deploy_demo,
],
outputs=[logs_box],
)
# Connect language change to phrase reloading and interface reveal (placed after all components are defined)
language_selector.change(
change_language,
inputs=[language_selector],
outputs=[phrase_texts_state, visible_rows_state] + phrase_markdowns + rec_components + [
add_rows_btn, record_dataset_btn, dataset_status, advanced_accordion,
save_rec_btn, push_recordings_btn, start_btn, logs_box
]
)
if __name__ == "__main__":
server_port = int(os.environ.get("INTERFACE_PORT", "7860"))
server_name = os.environ.get("INTERFACE_HOST", "0.0.0.0")
demo.queue().launch(server_name=server_name, server_port=server_port, mcp_server=True, ssr_mode=False)