Kyudan/whisperllama

inference code

import torch
from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq, AutoTokenizer, AutoModelForCausalLM
import soundfile as sf
from model import create_asr_model, modify_llama_blocks, decode_asr_output
import gc
import librosa
import numpy as np
import os
from datasets import load_dataset


def load_trained_model(model_path):
    gc.collect()
    torch.cuda.empty_cache()
    
    try:
        if torch.cuda.is_available():
            torch.cuda.set_per_process_memory_fraction(0.5)
        
        print("Loading Whisper encoder...")
        whisper = AutoModelForSpeechSeq2Seq.from_pretrained(
            "openai/whisper-large-v2",
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True,
            device_map="auto"  # 자동으로 메모리 관리
        )
        whisper_encoder = whisper.get_encoder()
        
        print("Loading Llama...")
        tokenizer = AutoTokenizer.from_pretrained(
            "meta-llama/Llama-3.2-1B",
            use_fast=True
        )
        
        # 토크나이저 설정
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.padding_side = "left"
        
        # Llama 모델 설정
        llama = AutoModelForCausalLM.from_pretrained(
            "meta-llama/Llama-3.2-1B",
            torch_dtype=torch.float16,
            low_cpu_mem_usage=True,
            device_map="auto"  # 자동으로 메모리 관리
        )
        llama.config.pad_token_id = tokenizer.pad_token_id
        llama.resize_token_embeddings(len(tokenizer))
        
        modified_llama = modify_llama_blocks(llama, num_blocks_to_keep=2)
        del llama
        gc.collect()
        
        print("Creating model...")
        model = create_asr_model(whisper_encoder, modified_llama)
        model = model.half()
        
        print("Loading weights...")
        state_dict = torch.load(model_path, map_location='cpu')
        
        # 디버깅 정보 출력
        print(f"\nModel vocab size: {model.decoder.model.embed_tokens.weight.shape[0]}")
        print(f"Tokenizer vocab size: {len(tokenizer)}")
        print(f"BOS token id: {tokenizer.bos_token_id}")
        print(f"EOS token id: {tokenizer.eos_token_id}")
        print(f"PAD token id: {tokenizer.pad_token_id}")
        
        missing, unexpected = model.load_state_dict(
            {k: v.half() for k, v in state_dict.items()}, 
            strict=False
        )
        
        print(f"\nMissing keys: {len(missing)}")
        print(f"Unexpected keys: {len(unexpected)}")
        processor = AutoProcessor.from_pretrained("openai/whisper-large-v2")

        model.eval()
        
        return model, processor, tokenizer
        
    except Exception as e:
        print(f"Error during model loading: {e}")
        gc.collect()
        torch.cuda.empty_cache()
        raise

def process_audio(audio_path, processor):
    try:
        print(f"Loading audio from {audio_path}...")
        # librosa를 사용하여 자동 리샘플링
        audio, orig_sr = librosa.load(audio_path)
        
        # 16kHz로 리샘플링
        if orig_sr != 16000:
            print(f"Resampling from {orig_sr}Hz to 16000Hz")
            audio = librosa.resample(audio, orig_sr=orig_sr, target_sr=16000)
        
        # 오디오 정규화
        audio = audio / np.abs(audio).max()
        
        input_features = processor(
            audio,
            sampling_rate=16000,
            return_tensors="pt"
        ).input_features.half()
        
        return input_features
        
    except Exception as e:
        print(f"Error processing audio: {e}")
        raise

def run_inference(model, input_features, tokenizer, max_length=200):
    try:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"\nUsing device: {device}")
        
        with torch.cuda.amp.autocast():
            with torch.no_grad():
                if torch.cuda.is_available():
                    model = model.to(device)
                    input_features = input_features.to(device)
                
                print("\nInput features shape:", input_features.shape)
                
                # 시작 토큰 설정
                start_token = tokenizer.bos_token_id
                print(f"Using start token: {start_token} ({tokenizer.decode([start_token])})")
                
                decoder_input_ids = torch.tensor([[start_token]], 
                                               device=device,
                                               dtype=torch.long)
                
                # Greedy decoding
                max_length = 100
                generated_ids = []
                
                for _ in range(max_length):
                    outputs = model(
                        input_features=input_features,
                        decoder_input_ids=decoder_input_ids
                    )
                    
                    next_token_logits = outputs.logits[:, -1, :]
                    next_token_id = torch.argmax(next_token_logits, dim=-1).item()
                    
                    # Top 5 토큰 출력
                    top_tokens = torch.topk(next_token_logits[0], 5)
                    print("\nTop 5 tokens for position", len(generated_ids))
                    for token_id, logit in zip(top_tokens.indices, top_tokens.values):
                        token = tokenizer.decode([token_id])
                        prob = torch.softmax(top_tokens.values, dim=-1)[0].item()
                        print(f"Token: {token}, Probability: {prob:.4f}")
                    
                    generated_ids.append(next_token_id)
                    
                    if next_token_id == tokenizer.eos_token_id:
                        break
                        
                    decoder_input_ids = torch.cat([
                        decoder_input_ids, 
                        torch.tensor([[next_token_id]], device=device)
                    ], dim=-1)
                
                # 전체 시퀀스 디코딩
                text = tokenizer.decode(generated_ids, skip_special_tokens=True)
                
                if torch.cuda.is_available():
                    model = model.cpu()
                    torch.cuda.empty_cache()
                
                return text
                
    except Exception as e:
        print(f"Error during inference: {e}")
        torch.cuda.empty_cache()
        raise
    finally:
        gc.collect()
        torch.cuda.empty_cache()


def main():
    try:
        model_path = "/home/elicer/.cache/huggingface/hub/models--Kyudan--whisperllama/snapshots/3269c93814c84e38f2d1a46935851f4923d73659/best_model_epoch_0.pt"
        
        # LibriSpeech 테스트 셋 로드 (10개 샘플)
        print("Loading LibriSpeech test samples...")
        dataset = load_dataset("librispeech_asr", "clean", split="test", streaming=True)
        samples = list(dataset.take(10))  # 10개 샘플만 가져오기
        
        print("Loading model...")
        model, processor, tokenizer = load_trained_model(model_path)
        
        # 각 샘플에 대해 추론 실행
        for idx, sample in enumerate(samples, 1):
            print(f"\nProcessing sample {idx}/10...")
            print(f"Speaker ID: {sample['speaker_id']}")
            print(f"Chapter ID: {sample['chapter_id']}")
            print(f"Reference text: {sample['text']}")
            
            # 오디오 처리
            input_features = processor(
                sample['audio']['array'],
                sampling_rate=16000,
                return_tensors="pt"
            ).input_features.half()
            
            # 추론 실행
            print("Running inference...")
            transcribed_text = run_inference(model, input_features, tokenizer)
            
            print("\nTranscription Results:")
            print("-" * 50)
            print(f"Model output: {transcribed_text}")
            print(f"Reference  : {sample['text']}")
            print("-" * 50)
            
    except Exception as e:
        print(f"Error in main: {e}")
    finally:
        gc.collect()
        torch.cuda.empty_cache()

if __name__ == "__main__":
    main()