추론+main 한꺼번에 작성, kobart interference 수정, whisper_v1 사용

app_v2/app_v2.py

@@ -0,0 +1,138 @@
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import JSONResponse
+from transformers import (
+    PreTrainedTokenizerFast,
+    BartForConditionalGeneration,
+    WhisperProcessor,
+    WhisperForConditionalGeneration,
+    AutoConfig
+)
+import torch
+import torchaudio
+import os
+import random
+
+app = FastAPI()
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# whisper 로딩
+whisper_model = WhisperForConditionalGeneration.from_pretrained(
+    "jeju-potato/jeju_potato_models", subfolder="whisper_v1_ep11"
+).to(device)
+whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-medium")
+whisper_model.config.forced_decoder_ids = None
+whisper_model.generation_config.forced_decoder_ids = None
+
+# first_kobart 로딩
+first_tokenizer = PreTrainedTokenizerFast.from_pretrained("jeju-potato/jeju_potato_models", subfolder="first_kobart_finetuned_v3")
+first_model = BartForConditionalGeneration.from_pretrained(
+    "jeju-potato/jeju_potato_models", subfolder="first_kobart_finetuned_v3"
+).to(device)
+first_model.config.forced_decoder_ids = None
+first_model.generation_config.forced_decoder_ids = None
+
+# second_kobart 로딩 (경고 제거용 config 명시)
+second_config = AutoConfig.from_pretrained(
+    "jeju-potato/jeju_potato_models", subfolder="second_kobart_finetuned_v2"
+)
+second_config.num_labels = 2
+second_config.id2label = {0: "NEGATIVE", 1: "POSITIVE"}
+second_config.label2id = {"NEGATIVE": 0, "POSITIVE": 1}
+
+second_tokenizer = PreTrainedTokenizerFast.from_pretrained("jeju-potato/jeju_potato_models", subfolder="second_kobart_finetuned_v2")
+second_model = BartForConditionalGeneration.from_pretrained(
+    "jeju-potato/jeju_potato_models",
+    subfolder="second_kobart_finetuned_v2",
+    config=second_config
+).to(device)
+second_model.config.forced_decoder_ids = None
+second_model.generation_config.forced_decoder_ids = None
+
+# whisper 추론
+def inference_whisper(audio_path):
+    waveform, sr = torchaudio.load(audio_path)
+    waveform = waveform.mean(dim=0) if waveform.shape[0] > 1 else waveform.squeeze(0)
+    inputs = whisper_processor(
+        waveform.numpy(),
+        sampling_rate=sr,
+        return_tensors="pt",
+        task="transcribe"
+    ).to(device)
+    whisper_model.generation_config.forced_decoder_ids = None
+    whisper_model.generation_config.suppress_tokens = []
+
+    with torch.no_grad():
+        pred_ids = whisper_model.generate(inputs["input_features"], max_length=128)
+        decoded = whisper_processor.batch_decode(pred_ids, skip_special_tokens=True)[0]
+    return decoded
+
+# kobart 추론
+def inference_kobart(text, model, tokenizer):
+    encoded = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
+    input_ids = encoded["input_ids"].to(device)
+    attention_mask = encoded["attention_mask"].to(device)
+
+    with torch.no_grad():
+        pred_ids = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_length=128
+        )
+        result = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)[0]
+    return result
+
+# 전체 파이프라인
+def run_pipeline(audio_path):
+    whisper_text = inference_whisper(audio_path)
+    first_kobart_text = inference_kobart(whisper_text, first_model, first_tokenizer)
+    second_kobart_text = inference_kobart(first_kobart_text, second_model, second_tokenizer)
+    return whisper_text, first_kobart_text, second_kobart_text
+
+# FastAPI 엔드포인트
+@app.get("/test_random")
+def test_random_file():
+    # data/source_data에서 랜덤으포 파일 하나 선택
+    audio_dir = r"D:\whisper\audio1_all"
+
+    candidates = [f for f in os.listdir(audio_dir) if f.endswith(".wav")]
+    if not candidates:
+        return {"error": "No .wav files found in data/source_data"}
+    
+    chosen_file = random.choice(candidates)
+    audio_path = os.path.join(audio_dir, chosen_file)
+
+    # 전체 파이프라인
+    try:
+        whisper_text, first_kobart_text, second_kobart_text = run_pipeline(audio_path)
+    except Exception as e:
+        return {"error": str(e)}
+
+    # 결과 반환
+    return {
+        "filename": chosen_file,
+        "whisper_result": whisper_text,
+        "first_kobart_result": first_kobart_text,
+        "second_kobart_result": second_kobart_text
+    }
+
+# 클라이언트로부터 음성파일을 Form-data 형식으로 받음
+@app.post("/inference")
+async def inference(audio: UploadFile=File(...)):
+    # 파일 저장
+    os.makedirs("temp", exist_ok=True)
+    audio_path = f"temp/{audio.filename}" # 클라이언트로부터 받은 음성파일을 temp/ 폴더에 저장
+    with open(audio_path, "wb") as f:
+        f.write(await audio.read())
+    
+    # 모델 추론
+    try:
+        whisper_text, first_kobart_text, second_kobart_text = run_pipeline(audio_path)
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})
+    
+    # 결과 반환
+    return {
+        "whisper_result": whisper_text,
+        "first_kobart_result": first_kobart_text,
+        "second_kobart_result": second_kobart_text
+    }