안드로이드 앱 연동, ngrok 주소 사용

app_v3/app_v3.py

@@ -0,0 +1,160 @@
+from prometheus_fastapi_instrumentator import Instrumentator
+from fastapi import FastAPI, UploadFile, File, Request, Response
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from transformers import (
+    PreTrainedTokenizerFast,
+    BartForConditionalGeneration,
+    WhisperProcessor,
+    WhisperForConditionalGeneration,
+    AutoConfig
+)
+import torch
+import torchaudio
+import os
+import random
+from pydub import AudioSegment
+from prometheus_client import Counter
+
+app = FastAPI()
+Instrumentator().instrument(app).expose(app)
+
+# ▶️ Custom Network Metrics
+request_bytes = Counter("request_bytes_total", "Total request bytes", ["path"])
+response_bytes = Counter("response_bytes_total", "Total response bytes", ["path"])
+
+@app.middleware("http")
+async def track_bytes(request: Request, call_next):
+    path = request.url.path
+    req_len = int(request.headers.get("content-length", 0))
+    request_bytes.labels(path=path).inc(req_len)
+
+    response = await call_next(request)
+
+    body = b""
+    async for chunk in response.body_iterator:
+        body += chunk
+    resp_len = len(body)
+
+    response_bytes.labels(path=path).inc(resp_len)
+    return Response(content=body, status_code=response.status_code, headers=dict(response.headers), media_type=response.media_type)
+
+# ▶️ Device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ▶️ CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ▶️ Whisper
+#local_model_path = r"D:\\whisper\\whisper_partial_ep16\\step30000-36000"
+#whisper_model = WhisperForConditionalGeneration.from_pretrained(local_model_path).to(device)
+whisper_model = WhisperForConditionalGeneration.from_pretrained(
+    "jeju-potato/jeju_potato_models", subfolder="whisper_v1_ep16"
+).to(device)
+whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-base")
+whisper_model.config.forced_decoder_ids = None
+whisper_model.generation_config.forced_decoder_ids = None
+
+# ▶️ KoBART 1
+first_tokenizer = PreTrainedTokenizerFast.from_pretrained("jeju-potato/jeju_potato_models", subfolder="first_kobart_v3")
+first_model = BartForConditionalGeneration.from_pretrained("jeju-potato/jeju_potato_models", subfolder="first_kobart_v3").to(device)
+first_model.config.forced_decoder_ids = None
+first_model.generation_config.forced_decoder_ids = None
+
+# ▶️ KoBART 2
+#second_config = AutoConfig.from_pretrained("jeju-potato/jeju_potato_models", subfolder="second_kobart_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_v2")
+second_model = BartForConditionalGeneration.from_pretrained("jeju-potato/jeju_potato_models", subfolder="second_kobart_v2").to(device)
+second_model.config.forced_decoder_ids = None
+second_model.generation_config.forced_decoder_ids = None
+
+# ▶️ Whisper inference
+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 inference
+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
+
+# ▶️ test_random
+def test_random_file():
+    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
+    }
+
+# ▶️ POST /inference
+@app.post("/inference")
+async def inference(audio: UploadFile = File(...)):
+    os.makedirs("temp", exist_ok=True)
+    input_path = f"temp/{audio.filename}"
+    with open(input_path, "wb") as f:
+        f.write(await audio.read())
+
+    output_path = "temp/converted.wav"
+    try:
+        sound = AudioSegment.from_file(input_path)
+        sound = sound.set_frame_rate(16000).set_channels(1)
+        sound.export(output_path, format="wav")
+    except Exception as e:
+        return JSONResponse(status_code=400, content={"error": f"와이파 변환 실패: {e}"})
+
+    try:
+        whisper_text, first_kobart_text, second_kobart_text = run_pipeline(output_path)
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": f"모델 추론 실패: {e}"})
+
+    return {
+        "whisper_result": whisper_text,
+        "first_kobart_result": first_kobart_text,
+        "second_kobart_result": second_kobart_text
+    }