Sure! Pl
This commit is contained in:
Alireza
190
vosk/test_files/batch_confirm_hf.py
Normal file
190
vosk/test_files/batch_confirm_hf.py
Normal file
@@ -0,0 +1,190 @@
|
||||
from datasets import load_dataset, Audio, Dataset
|
||||
import soundfile as sf
|
||||
import requests
|
||||
import os
|
||||
from tqdm import tqdm
|
||||
import pandas as pd
|
||||
import json
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
import numpy as np
|
||||
from huggingface_hub import HfApi, create_repo
|
||||
|
||||
# Load the dataset with audio decoding
|
||||
print("Loading dataset...")
|
||||
ds = load_dataset(
|
||||
"Ashegh-Sad-Warrior/Persian_Common_Voice_17_0",
|
||||
split="validated[:500]",
|
||||
streaming=False
|
||||
).cast_column("audio", Audio(sampling_rate=16000))
|
||||
|
||||
output_dir = "confirmed_dataset"
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
confirmed = []
|
||||
|
||||
API_URL = "http://localhost:5000/batch_confirm"
|
||||
batch_size = 8
|
||||
|
||||
# Hugging Face configuration
|
||||
HF_DATASET_NAME = "dpr2000/persian-cv17-confirmed" # Change this to your desired dataset name
|
||||
HF_PRIVATE = True # Set to True for private dataset
|
||||
|
||||
def save_flac(audio_array, path):
|
||||
sf.write(path, audio_array, 16000, format="FLAC")
|
||||
|
||||
print("Processing batches...")
|
||||
for i in tqdm(range(0, len(ds), batch_size)):
|
||||
batch = ds[i:i+batch_size]
|
||||
files = {}
|
||||
references = []
|
||||
temp_flacs = []
|
||||
audio_arrays = []
|
||||
# Fix: batch is a dict of lists
|
||||
for j in range(len(batch["audio"])):
|
||||
audio = batch["audio"][j]
|
||||
flac_path = f"temp_{i+j}.flac"
|
||||
save_flac(audio["array"], flac_path)
|
||||
files[f"audio{j}"] = open(flac_path, "rb")
|
||||
references.append(batch["sentence"][j])
|
||||
temp_flacs.append(flac_path)
|
||||
audio_arrays.append(audio["array"]) # Store the array for confirmed
|
||||
data = {"references": json.dumps(references)}
|
||||
try:
|
||||
response = requests.post(API_URL, files=files, data=data, timeout=120)
|
||||
if response.status_code == 200:
|
||||
resp_json = response.json()
|
||||
if "results" in resp_json:
|
||||
results = resp_json["results"]
|
||||
else:
|
||||
print(f"Batch {i} failed: 'results' key missing in response: {resp_json}")
|
||||
results = [None] * len(references)
|
||||
else:
|
||||
print(f"Batch {i} failed: HTTP {response.status_code} - {response.text}")
|
||||
results = [None] * len(references)
|
||||
except Exception as e:
|
||||
print(f"Batch {i} failed: {e}")
|
||||
results = [None] * len(references)
|
||||
for j, result in enumerate(results):
|
||||
if result and result.get("confirmed"):
|
||||
# Save confirmed audio array and transcription
|
||||
confirmed.append({"audio": audio_arrays[j], "transcription": references[j]})
|
||||
os.remove(temp_flacs[j])
|
||||
else:
|
||||
os.remove(temp_flacs[j])
|
||||
for f in files.values():
|
||||
f.close()
|
||||
|
||||
# Save confirmed data using sharding approach
|
||||
if confirmed:
|
||||
print(f"\n🔄 Saving {len(confirmed)} confirmed samples...")
|
||||
|
||||
# Convert confirmed data to HuggingFace dataset format
|
||||
def extract_minimal(example):
|
||||
# Convert float32 audio (range -1.0 to 1.0) to int16 (range -32768 to 32767)
|
||||
audio_float32 = np.array(example["audio"], dtype=np.float32)
|
||||
# Ensure audio is in valid range and scale to int16
|
||||
audio_float32 = np.clip(audio_float32, -1.0, 1.0)
|
||||
audio_int16 = (audio_float32 * 32767).astype(np.int16)
|
||||
return {
|
||||
"audio": audio_int16.tobytes(), # Store as int16 bytes, compatible with Whisper
|
||||
"text": example["transcription"]
|
||||
}
|
||||
|
||||
# Create dataset from confirmed samples
|
||||
confirmed_dataset = Dataset.from_list(confirmed)
|
||||
confirmed_dataset = confirmed_dataset.map(extract_minimal, remove_columns=confirmed_dataset.column_names)
|
||||
|
||||
# Sharding parameters
|
||||
num_shards = min(1, len(confirmed)) # Don't create more shards than samples
|
||||
shard_size = len(confirmed_dataset) // num_shards + 1
|
||||
|
||||
# Write each shard separately
|
||||
for i in range(num_shards):
|
||||
start = i * shard_size
|
||||
end = min(len(confirmed_dataset), (i + 1) * shard_size)
|
||||
|
||||
if start >= len(confirmed_dataset):
|
||||
break
|
||||
|
||||
shard = confirmed_dataset.select(range(start, end))
|
||||
table = pa.Table.from_pandas(shard.to_pandas()) # Convert to PyArrow table
|
||||
|
||||
shard_path = os.path.join(output_dir, f"confirmed_shard_{i:02}.parquet")
|
||||
|
||||
pq.write_table(
|
||||
table,
|
||||
shard_path,
|
||||
compression="zstd",
|
||||
compression_level=22, # Maximum compression
|
||||
use_dictionary=True,
|
||||
version="2.6"
|
||||
)
|
||||
|
||||
print(f"🔹 Shard {i+1}/{num_shards}: {len(shard)} samples saved")
|
||||
|
||||
print(f"\n✅ All confirmed data saved in {num_shards} shards in `{output_dir}/`")
|
||||
|
||||
# Push to Hugging Face Hub
|
||||
print(f"\n🚀 Pushing dataset to Hugging Face Hub as '{HF_DATASET_NAME}'...")
|
||||
try:
|
||||
# Initialize HF API
|
||||
api = HfApi()
|
||||
|
||||
# Create the repository (private if specified)
|
||||
try:
|
||||
create_repo(
|
||||
repo_id=HF_DATASET_NAME,
|
||||
repo_type="dataset",
|
||||
private=HF_PRIVATE,
|
||||
exist_ok=True
|
||||
)
|
||||
print(f"✅ Repository '{HF_DATASET_NAME}' created/verified")
|
||||
except Exception as e:
|
||||
print(f"⚠️ Repository creation: {e}")
|
||||
|
||||
# Upload all parquet files
|
||||
for i in range(num_shards):
|
||||
shard_path = os.path.join(output_dir, f"confirmed_shard_{i:02}.parquet")
|
||||
if os.path.exists(shard_path):
|
||||
api.upload_file(
|
||||
path_or_fileobj=shard_path,
|
||||
path_in_repo=f"confirmed_shard_{i:02}.parquet",
|
||||
repo_id=HF_DATASET_NAME,
|
||||
repo_type="dataset"
|
||||
)
|
||||
print(f"📤 Uploaded shard {i+1}/{num_shards}")
|
||||
|
||||
# Create dataset info file
|
||||
dataset_info = {
|
||||
"dataset_name": HF_DATASET_NAME,
|
||||
"description": "Persian Common Voice confirmed samples for Whisper fine-tuning",
|
||||
"total_samples": len(confirmed),
|
||||
"num_shards": num_shards,
|
||||
"audio_format": "int16 PCM, 16kHz",
|
||||
"columns": ["audio", "text"],
|
||||
"source_dataset": "Ashegh-Sad-Warrior/Persian_Common_Voice_17_0",
|
||||
"processing": "Vosk API batch confirmation"
|
||||
}
|
||||
|
||||
# Upload dataset info
|
||||
import tempfile
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) as f:
|
||||
json.dump(dataset_info, f, indent=2, ensure_ascii=False)
|
||||
info_path = f.name
|
||||
|
||||
api.upload_file(
|
||||
path_or_fileobj=info_path,
|
||||
path_in_repo="dataset_info.json",
|
||||
repo_id=HF_DATASET_NAME,
|
||||
repo_type="dataset"
|
||||
)
|
||||
os.unlink(info_path)
|
||||
|
||||
print(f"🎉 Dataset successfully pushed to: https://huggingface.co/datasets/{HF_DATASET_NAME}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Failed to push to Hugging Face: {e}")
|
||||
print("💡 Make sure you're logged in with: huggingface-cli login")
|
||||
|
||||
else:
|
||||
print("❌ No confirmed samples to save")
|
||||
Reference in New Issue
Block a user