update README for finetuning-asr
This commit is contained in:
YaoyaoChang
@@ -0,0 +1,234 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
Inference with LoRA Fine-tuned VibeVoice ASR Model
|
||||
|
||||
This script loads a LoRA fine-tuned model and runs inference.
|
||||
|
||||
Usage:
|
||||
python inference_lora.py \
|
||||
--base_model microsoft/VibeVoice-ASR \
|
||||
--lora_path ./output \
|
||||
--audio_file ./toy_dataset/0.mp3
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import torch
|
||||
|
||||
from peft import PeftModel
|
||||
|
||||
from vibevoice.modular.modeling_vibevoice_asr import VibeVoiceASRForConditionalGeneration
|
||||
from vibevoice.processor.vibevoice_asr_processor import VibeVoiceASRProcessor
|
||||
|
||||
|
||||
def load_lora_model(
|
||||
base_model_path: str,
|
||||
lora_path: str,
|
||||
device: str = "cuda",
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
):
|
||||
"""
|
||||
Load base model and merge with LoRA weights.
|
||||
|
||||
Args:
|
||||
base_model_path: Path to base pretrained model
|
||||
lora_path: Path to LoRA adapter weights
|
||||
device: Device to load model on
|
||||
dtype: Data type for model
|
||||
|
||||
Returns:
|
||||
Tuple of (model, processor)
|
||||
"""
|
||||
print(f"Loading base model from {base_model_path}")
|
||||
|
||||
# Load processor
|
||||
processor = VibeVoiceASRProcessor.from_pretrained(
|
||||
base_model_path,
|
||||
language_model_pretrained_name="Qwen/Qwen2.5-7B"
|
||||
)
|
||||
|
||||
# Load base model
|
||||
model = VibeVoiceASRForConditionalGeneration.from_pretrained(
|
||||
base_model_path,
|
||||
dtype=dtype,
|
||||
device_map=device if device == "auto" else None,
|
||||
attn_implementation="flash_attention_2",
|
||||
trust_remote_code=True,
|
||||
)
|
||||
|
||||
if device != "auto":
|
||||
model = model.to(device)
|
||||
|
||||
# Load LoRA adapter
|
||||
print(f"Loading LoRA adapter from {lora_path}")
|
||||
model = PeftModel.from_pretrained(model, lora_path)
|
||||
|
||||
# Optionally merge LoRA weights into base model for faster inference
|
||||
# model = model.merge_and_unload()
|
||||
|
||||
model.eval()
|
||||
print("Model loaded successfully")
|
||||
|
||||
return model, processor
|
||||
|
||||
|
||||
def transcribe(
|
||||
model,
|
||||
processor,
|
||||
audio_path: str,
|
||||
max_new_tokens: int = 4096,
|
||||
temperature: float = 0.0,
|
||||
context_info: str = None,
|
||||
device: str = "cuda",
|
||||
):
|
||||
"""
|
||||
Transcribe an audio file using the LoRA fine-tuned model.
|
||||
|
||||
Args:
|
||||
model: The LoRA fine-tuned model
|
||||
processor: The processor
|
||||
audio_path: Path to audio file
|
||||
max_new_tokens: Maximum tokens to generate
|
||||
temperature: Sampling temperature (0 = greedy)
|
||||
context_info: Optional context info (e.g., hotwords)
|
||||
device: Device
|
||||
|
||||
Returns:
|
||||
Transcription result
|
||||
"""
|
||||
print(f"\nTranscribing: {audio_path}")
|
||||
|
||||
# Process audio
|
||||
inputs = processor(
|
||||
audio=audio_path,
|
||||
sampling_rate=None,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
add_generation_prompt=True,
|
||||
context_info=context_info,
|
||||
)
|
||||
|
||||
# Move to device
|
||||
inputs = {k: v.to(device) if isinstance(v, torch.Tensor) else v
|
||||
for k, v in inputs.items()}
|
||||
|
||||
# Generation config
|
||||
gen_config = {
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"pad_token_id": processor.pad_id,
|
||||
"eos_token_id": processor.tokenizer.eos_token_id,
|
||||
"do_sample": temperature > 0,
|
||||
}
|
||||
if temperature > 0:
|
||||
gen_config["temperature"] = temperature
|
||||
gen_config["top_p"] = 0.9
|
||||
|
||||
# Generate
|
||||
with torch.no_grad():
|
||||
output_ids = model.generate(**inputs, **gen_config)
|
||||
|
||||
# Decode
|
||||
input_length = inputs['input_ids'].shape[1]
|
||||
generated_ids = output_ids[0, input_length:]
|
||||
generated_text = processor.decode(generated_ids, skip_special_tokens=True)
|
||||
|
||||
# Parse structured output
|
||||
try:
|
||||
segments = processor.post_process_transcription(generated_text)
|
||||
except Exception as e:
|
||||
print(f"Warning: Failed to parse structured output: {e}")
|
||||
segments = []
|
||||
|
||||
return {
|
||||
"raw_text": generated_text,
|
||||
"segments": segments,
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Inference with LoRA Fine-tuned VibeVoice ASR")
|
||||
|
||||
parser.add_argument(
|
||||
"--base_model",
|
||||
type=str,
|
||||
default="microsoft/VibeVoice-ASR",
|
||||
help="Path to base pretrained model"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lora_path",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to LoRA adapter weights"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio_file",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to audio file to transcribe"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--context_info",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Optional context info (e.g., 'Hotwords: Tea Brew, Aiden Host')"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max_new_tokens",
|
||||
type=int,
|
||||
default=4096,
|
||||
help="Maximum tokens to generate"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--temperature",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="Sampling temperature (0 = greedy)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--device",
|
||||
type=str,
|
||||
default="cuda" if torch.cuda.is_available() else "cpu",
|
||||
help="Device to use"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Load model
|
||||
dtype = torch.bfloat16 if args.device != "cpu" else torch.float32
|
||||
model, processor = load_lora_model(
|
||||
base_model_path=args.base_model,
|
||||
lora_path=args.lora_path,
|
||||
device=args.device,
|
||||
dtype=dtype,
|
||||
)
|
||||
|
||||
# Transcribe
|
||||
result = transcribe(
|
||||
model=model,
|
||||
processor=processor,
|
||||
audio_path=args.audio_file,
|
||||
max_new_tokens=args.max_new_tokens,
|
||||
temperature=args.temperature,
|
||||
context_info=args.context_info,
|
||||
device=args.device,
|
||||
)
|
||||
|
||||
# Print results
|
||||
print("\n" + "="*60)
|
||||
print("Transcription Result")
|
||||
print("="*60)
|
||||
|
||||
print("\n--- Raw Output ---")
|
||||
raw_text = result['raw_text']
|
||||
print(raw_text[:2000] + "..." if len(raw_text) > 2000 else raw_text)
|
||||
|
||||
if result['segments']:
|
||||
print(f"\n--- Structured Output ({len(result['segments'])} segments) ---")
|
||||
for seg in result['segments'][:20]:
|
||||
print(f"[{seg.get('start_time', 'N/A')} - {seg.get('end_time', 'N/A')}] "
|
||||
f"Speaker {seg.get('speaker_id', 'N/A')}: {seg.get('text', '')[:80]}...")
|
||||
if len(result['segments']) > 20:
|
||||
print(f" ... and {len(result['segments']) - 20} more segments")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user