app.py

import os, asyncio, json, tempfile, websockets, pdfplumber
import gradio as gr
import openai
from dotenv import load_dotenv
import numpy as np
import wave

# ─── 0. 초기화 ───────────────────────────────────────────────
load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")
if not openai.api_key:
    raise RuntimeError("OPENAI_API_KEY 가 .env 에 없습니다!")

LANG = ["Korean","English","Japanese","Chinese",
        "Thai","Russian","Vietnamese","Spanish","French"]
VOICE = {l: ("nova" if l in ["Korean","Japanese","Chinese"] else "alloy")
         for l in LANG}
FOUR = ["English","Chinese","Thai","Russian"]
WS_URL = "wss://api.openai.com/v1/realtime"  # 올바른 엔드포인트로 수정

# ─── 1. 공통 GPT 번역 / TTS ─────────────────────────────────
async def gpt_translate(text, src, tgt):
    rsp = await openai.AsyncClient().chat.completions.create(
        model="gpt-3.5-turbo",
        messages=[{"role":"system",
                   "content":f"Translate {src} → {tgt}. Return only the text."},
                  {"role":"user","content":text}],
        temperature=0.3,max_tokens=2048)
    return rsp.choices[0].message.content.strip()

async def gpt_tts(text, lang):
    rsp = await openai.AsyncClient().audio.speech.create(
        model="tts-1", voice=VOICE[lang], input=text[:4096])
    tmp = tempfile.NamedTemporaryFile(delete=False,suffix=".mp3")
    tmp.write(rsp.content); tmp.close(); return tmp.name

# ─── 2. PDF 번역 ────────────────────────────────────────────
def translate_pdf(file, src, tgt):
    if not file: return "⚠️ PDF 업로드 필요", ""
    with pdfplumber.open(file.name) as pdf:
        text = "\n".join(p.extract_text() or "" for p in pdf.pages[:5]).strip()
    if not text:
        return "⚠️ 텍스트 추출 실패", ""
    return text, asyncio.run(gpt_translate(text, src, tgt))

# ─── 2-1. 오디오 번역 (탭1용) ────────────────────────────────
async def translate_audio_async(file, src, tgt):
    if not file: return "⚠️ 오디오 업로드 필요", "", None
    
    # STT: Whisper API 사용
    with open(file, 'rb') as audio_file:
        transcript = await openai.AsyncClient().audio.transcriptions.create(
            model="whisper-1",
            file=audio_file,
            language=src[:2].lower()  # 언어 코드 간소화
        )
    
    orig_text = transcript.text
    trans_text = await gpt_translate(orig_text, src, tgt)
    audio_path = await gpt_tts(trans_text, tgt)
    
    return orig_text, trans_text, audio_path

def translate_audio(file, src, tgt):
    return asyncio.run(translate_audio_async(file, src, tgt))

# ─── 3. 실시간 STT (Whisper API 사용) ──────────────────────────
async def process_audio_chunk(audio_data, src_lang):
    """오디오 청크를 처리하여 텍스트로 변환"""
    if audio_data is None:
        return ""
    
    try:
        # Gradio는 (sample_rate, audio_array) 튜플을 반환
        if isinstance(audio_data, tuple):
            sample_rate, audio_array = audio_data
            
            # 오디오가 너무 짧으면 무시 (0.5초 미만)
            if len(audio_array) < sample_rate * 0.5:
                return ""
            
            # 오디오 정규화 및 노이즈 필터링
            audio_array = audio_array.astype(np.float32)
            
            # 무음 감지 - RMS가 너무 낮으면 무시
            rms = np.sqrt(np.mean(audio_array**2))
            if rms < 0.01:  # 무음 임계값
                return ""
            
            # 정규화
            max_val = np.max(np.abs(audio_array))
            if max_val > 0:
                audio_array = audio_array / max_val * 0.95
            
            # numpy array를 WAV 파일로 변환
            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
                with wave.open(tmp.name, 'wb') as wav_file:
                    wav_file.setnchannels(1)  # mono
                    wav_file.setsampwidth(2)  # 16-bit
                    wav_file.setframerate(sample_rate)
                    
                    # float32를 16-bit PCM으로 변환
                    audio_int16 = (audio_array * 32767).astype(np.int16)
                    wav_file.writeframes(audio_int16.tobytes())
                tmp_path = tmp.name
        else:
            # bytes 데이터인 경우
            with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
                tmp.write(audio_data)
                tmp_path = tmp.name
        
        # Whisper API로 변환 - 언어 힌트와 프롬프트 추가
        with open(tmp_path, 'rb') as audio_file:
            # 언어별 프롬프트 설정으로 hallucination 방지
            language_prompts = {
                "Korean": "이것은 한국어 대화입니다.",
                "English": "This is an English conversation.",
                "Japanese": "これは日本語の会話です。",
                "Chinese": "这是中文对话。",
            }
            
            prompt = language_prompts.get(src_lang, "")
            
            transcript = await openai.AsyncClient().audio.transcriptions.create(
                model="whisper-1",
                file=audio_file,
                language=src_lang[:2].lower(),
                prompt=prompt,
                temperature=0.0  # 더 보수적인 추론
            )
        
        os.unlink(tmp_path)  # 임시 파일 삭제
        
        # 결과 후처리 - 반복되는 패턴 제거
        text = transcript.text.strip()
        
        # 같은 문장이 반복되는 경우 처리
        sentences = text.split('.')
        if len(sentences) > 1:
            unique_sentences = []
            for sent in sentences:
                sent = sent.strip()
                if sent and (not unique_sentences or sent != unique_sentences[-1]):
                    unique_sentences.append(sent)
            text = '. '.join(unique_sentences)
            if text and not text.endswith('.'):
                text += '.'
        
        # 뉴스 관련 hallucination 패턴 감지 및 제거
        hallucination_patterns = [
            "MBC 뉴스", "KBS 뉴스", "SBS 뉴스", "JTBC 뉴스", 
            "뉴스룸", "뉴스데스크", "앵커", "기자입니다"
        ]
        
        # 짧은 텍스트에서 뉴스 패턴이 감지되면 무시
        if len(text) < 50 and any(pattern in text for pattern in hallucination_patterns):
            return ""
        
        return text
        
    except Exception as e:
        print(f"STT 오류: {e}")
        return ""

# ─── 4. Gradio 스트림 핸들러 (동기 버전) ─────────────────────
def realtime_single_sync(audio, src, tgt, state):
    """동기 버전의 실시간 단일 언어 번역"""
    if state is None:
        state = {"orig": "", "trans": "", "audio_buffer": [], "sample_rate": None}
    
    if audio is None:
        # 스트림 종료 시 남은 버퍼 처리
        if state["audio_buffer"] and state["sample_rate"]:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                # 버퍼의 오디오 합치기
                combined_audio = np.concatenate(state["audio_buffer"])
                audio_data = (state["sample_rate"], combined_audio)
                
                text = loop.run_until_complete(process_audio_chunk(audio_data, src))
                if text:
                    state["orig"] = state["orig"] + " " + text if state["orig"] else text
                    trans = loop.run_until_complete(gpt_translate(text, src, tgt))
                    state["trans"] = state["trans"] + " " + trans if state["trans"] else trans
            finally:
                loop.close()
            state["audio_buffer"] = []
        
        return state["orig"], state["trans"], state
    
    # 오디오 데이터 버퍼링
    if isinstance(audio, tuple):
        sample_rate, audio_array = audio
        state["sample_rate"] = sample_rate
        state["audio_buffer"].append(audio_array)
        
        # 버퍼가 충분히 쌓였을 때만 처리 (약 2-3초 분량)
        buffer_duration = len(np.concatenate(state["audio_buffer"])) / sample_rate
        if buffer_duration >= 2.0:  # 2초마다 처리
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            
            try:
                # 버퍼의 오디오 합치기
                combined_audio = np.concatenate(state["audio_buffer"])
                audio_data = (sample_rate, combined_audio)
                
                # STT
                text = loop.run_until_complete(process_audio_chunk(audio_data, src))
                if text:
                    state["orig"] = state["orig"] + " " + text if state["orig"] else text
                    
                    # 번역
                    trans = loop.run_until_complete(gpt_translate(text, src, tgt))
                    state["trans"] = state["trans"] + " " + trans if state["trans"] else trans
                
                # 버퍼 초기화
                state["audio_buffer"] = []
            finally:
                loop.close()
    
    return state["orig"], state["trans"], state

def realtime_four_sync(audio, src, state):
    """동기 버전의 실시간 4언어 번역"""
    if state is None:
        state = {"orig": "", "English": "", "Chinese": "", "Thai": "", "Russian": "", 
                 "audio_buffer": [], "sample_rate": None}
    
    if audio is None:
        # 스트림 종료 시 남은 버퍼 처리
        if state["audio_buffer"] and state["sample_rate"]:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                combined_audio = np.concatenate(state["audio_buffer"])
                audio_data = (state["sample_rate"], combined_audio)
                
                text = loop.run_until_complete(process_audio_chunk(audio_data, src))
                if text:
                    state["orig"] = state["orig"] + " " + text if state["orig"] else text
                    
                    tasks = []
                    for lang in FOUR:
                        tasks.append(gpt_translate(text, src, lang))
                    
                    translations = loop.run_until_complete(asyncio.gather(*tasks))
                    
                    for lang, trans in zip(FOUR, translations):
                        state[lang] = state[lang] + " " + trans if state[lang] else trans
            finally:
                loop.close()
            state["audio_buffer"] = []
        
        return (state["orig"], state["English"], state["Chinese"], 
                state["Thai"], state["Russian"], state)
    
    # 오디오 데이터 버퍼링
    if isinstance(audio, tuple):
        sample_rate, audio_array = audio
        state["sample_rate"] = sample_rate
        state["audio_buffer"].append(audio_array)
        
        # 버퍼가 충분히 쌓였을 때만 처리
        buffer_duration = len(np.concatenate(state["audio_buffer"])) / sample_rate
        if buffer_duration >= 2.0:  # 2초마다 처리
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            
            try:
                combined_audio = np.concatenate(state["audio_buffer"])
                audio_data = (sample_rate, combined_audio)
                
                # STT
                text = loop.run_until_complete(process_audio_chunk(audio_data, src))
                if text:
                    state["orig"] = state["orig"] + " " + text if state["orig"] else text
                    
                    # 4개 언어로 번역
                    tasks = []
                    for lang in FOUR:
                        tasks.append(gpt_translate(text, src, lang))
                    
                    translations = loop.run_until_complete(asyncio.gather(*tasks))
                    
                    for lang, trans in zip(FOUR, translations):
                        state[lang] = state[lang] + " " + trans if state[lang] else trans
                
                state["audio_buffer"] = []
            finally:
                loop.close()
    
    return (state["orig"], state["English"], state["Chinese"], 
            state["Thai"], state["Russian"], state)

# ─── 5. UI ──────────────────────────────────────────────────
with gr.Blocks(title="SMARTok Demo") as demo:
    with gr.Tabs():
        # 탭 1 – 오디오 번역
        with gr.TabItem("🎙️ 오디오"):
            src1 = gr.Dropdown(LANG, value="Korean", label="입력 언어")
            tgt1 = gr.Dropdown(LANG, value="English", label="출력 언어")
            aud1 = gr.Audio(sources=["microphone", "upload"], type="filepath")
            btn1 = gr.Button("번역")
            o1 = gr.Textbox(label="원문")
            t1 = gr.Textbox(label="번역")
            a1 = gr.Audio(label="TTS", type="filepath", autoplay=True)
            
            btn1.click(translate_audio, [aud1, src1, tgt1], [o1, t1, a1])

        # 탭 2 – PDF 번역
        with gr.TabItem("📄 PDF"):
            src2 = gr.Dropdown(LANG, value="Korean", label="입력 언어")
            tgt2 = gr.Dropdown(LANG, value="English", label="출력 언어")
            pdf = gr.File(file_types=[".pdf"])
            btn2 = gr.Button("번역")
            o2 = gr.Textbox(label="추출 원문", lines=15)
            t2 = gr.Textbox(label="번역 결과", lines=15)
            
            btn2.click(translate_pdf, [pdf, src2, tgt2], [o2, t2])

        # 탭 3 – 실시간 1언어
        with gr.TabItem("⏱️ 실시간 1"):
            src3 = gr.Dropdown(LANG, value="Korean", label="입력 언어")
            tgt3 = gr.Dropdown(LANG, value="English", label="출력 언어")
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("🎤 **마이크 입력**")
                    mic3 = gr.Audio(
                        sources=["microphone"], 
                        streaming=True,
                        type="numpy",  # numpy 형식 명시
                        label="마이크"
                    )
                    gr.Markdown("💡 **사용 방법**\n- 2-3초 정도 문장을 말씀해주세요\n- 너무 짧거나 긴 문장은 인식이 어려울 수 있습니다")
                
                with gr.Column():
                    o3 = gr.Textbox(label="원문(실시간)", lines=8, interactive=False)
                    t3 = gr.Textbox(label="번역(실시간)", lines=8, interactive=False)
            
            st3 = gr.State()
            
            # stream 메서드 수정
            mic3.stream(
                realtime_single_sync,
                inputs=[mic3, src3, tgt3, st3],
                outputs=[o3, t3, st3],
                time_limit=30,  # 30초 제한
                stream_every=0.5  # 0.5초마다 스트림
            )

        # 탭 4 – 실시간 4언어
        with gr.TabItem("🌏 실시간 4"):
            src4 = gr.Dropdown(LANG, value="Korean", label="입력 언어")
            
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("🎤 **마이크 입력**")
                    mic4 = gr.Audio(
                        sources=["microphone"], 
                        streaming=True,
                        type="numpy",
                        label="마이크"
                    )
                    o4 = gr.Textbox(label="원문", lines=8, interactive=False)
                
                with gr.Column(scale=2):
                    with gr.Row():
                        e4 = gr.Textbox(label="English", lines=8, interactive=False)
                        c4 = gr.Textbox(label="Chinese(简体)", lines=8, interactive=False)
                    with gr.Row():
                        th4 = gr.Textbox(label="Thai", lines=8, interactive=False)
                        r4 = gr.Textbox(label="Russian", lines=8, interactive=False)
            
            st4 = gr.State()
            
            # stream 메서드 수정
            mic4.stream(
                realtime_four_sync,
                inputs=[mic4, src4, st4],
                outputs=[o4, e4, c4, th4, r4, st4],
                time_limit=30,
                stream_every=0.5
            )

demo.launch(server_name="0.0.0.0", server_port=7860, debug=True)