app.py

# app.py
import streamlit as st
from faster_whisper import WhisperModel
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import tempfile
import os
import base64
import warnings

warnings.filterwarnings("ignore")

st.set_page_config(
    page_title="🎙️ Audio Transcriber + Chat IA Local",
    page_icon="🎙️",
    layout="wide"
)

# ============================================
# CARGA DE MODELOS
# ============================================

@st.cache_resource
def load_stt():
    return WhisperModel("small", device="cpu", compute_type="int8")

@st.cache_resource
def load_llm():
    model_name = "microsoft/phi-2"
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        dtype=torch.float32,
        device_map="auto",
        trust_remote_code=True
    )
    return model, tokenizer

# ============================================
# FUNCIONES
# ============================================

def transcribe_audio(audio_path, language=None):
    model = load_stt()
    segments_result = []
    
    segments, info = model.transcribe(
        audio_path,
        language=language,
        beam_size=5,
        word_timestamps=True,
        vad_filter=True
    )
    
    for seg in segments:
        segments_result.append({
            "text": seg.text.strip(),
            "start": seg.start,
            "end": seg.end
        })
    
    detected_lang = language if language else info.language
    return segments_result, detected_lang

def chat_with_context(question, segments, model, tokenizer):
    context_parts = []
    for seg in segments[:10]:
        context_parts.append(f"[{seg['start']:.1f}s] {seg['text']}")
    context = "\n".join(context_parts)
    
    spanish_chars = set('áéíóúñ¿¡')
    is_spanish = any(c in question.lower() for c in spanish_chars)
    
    if is_spanish:
        prompt = f"""Contexto de la transcripción:
{context}

Pregunta: {question}

Responde basándote SOLO en el contexto proporcionado. Si la respuesta no está en el contexto, di "No encuentro esa información en la transcripción".
Respuesta:"""
    else:
        prompt = f"""Transcript context:
{context}

Question: {question}

Answer based ONLY on the context provided. If the answer is not in the context, say "I cannot find that information in the transcript".
Answer:"""
    
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=300,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id
        )
    
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    if "Respuesta:" in answer:
        answer = answer.split("Respuesta:")[-1].strip()
    elif "Answer:" in answer:
        answer = answer.split("Answer:")[-1].strip()
    
    return answer

def get_audio_base64(audio_path):
    """Convierte audio a base64 para incrustar en HTML"""
    with open(audio_path, "rb") as f:
        audio_bytes = f.read()
    return base64.b64encode(audio_bytes).decode()

# ============================================
# INTERFAZ PRINCIPAL
# ============================================

def main():
    st.title("🎙️ Transcripción de Audio + Chat con IA Local")
    st.markdown("*Reproducción sincronizada con autoscroll*")
    
    with st.sidebar:
        st.header("📁 Configuración")
        
        audio_file = st.file_uploader(
            "Subir archivo de audio",
            type=["mp3", "wav", "m4a", "ogg", "flac"]
        )
        
        lang = st.selectbox(
            "Idioma del audio",
            ["Auto-detectado", "Español", "English"]
        )
        
        lang_map = {
            "Auto-detectado": None,
            "Español": "es",
            "English": "en"
        }
        
        if audio_file is not None:
            if st.button("🎯 Transcribir Audio", type="primary", use_container_width=True):
                # Guardar archivo temporal
                file_extension = os.path.splitext(audio_file.name)[1]
                with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as tmp:
                    tmp.write(audio_file.getbuffer())
                    tmp_path = tmp.name
                
                with st.spinner("Transcribiendo audio..."):
                    segments, detected_lang = transcribe_audio(tmp_path, lang_map[lang])
                
                st.session_state.segments = segments
                st.session_state.detected_lang = detected_lang
                st.session_state.audio_path = tmp_path
                st.session_state.full_text = " ".join([s["text"] for s in segments])
                
                st.success(f"✅ Transcripción completada - Idioma: {detected_lang.upper()}")
                st.rerun()
        
        with st.expander("ℹ️ Información"):
            st.caption("🎤 faster-whisper (small)")
            st.caption("🧠 microsoft/phi-2 (2.7B)")
            st.caption("🔊 Audio sincronizado con texto")
    
    # ========== PANEL PRINCIPAL ==========
    if "segments" in st.session_state and "audio_path" in st.session_state:
        col1, col2 = st.columns([2, 1])
        
        with col1:
            st.subheader("🎧 Reproductor de Audio")
            
            # Convertir audio a base64 para incrustar
            audio_base64 = get_audio_base64(st.session_state.audio_path)
            audio_mime = "audio/mpeg" if st.session_state.audio_path.endswith(('.mp3', '.m4a')) else "audio/wav"
            
            # Integrar reproductor y transcripción en un solo HTML para que el JS funcione
            import json
            segments_json = json.dumps([
                {"start": seg["start"], "end": seg["end"]} for seg in st.session_state.segments
            ])
            
            segments_html = ""
            for i, seg in enumerate(st.session_state.segments):
                segments_html += f"""
                <div class="segment" id="segment_{i}" onclick="jumpTo({seg['start']})">
                    <div class="time">{seg['start']:.1f}s</div>
                    <div class="text">{seg['text']}</div>
                    <button class="btn" title="Ir a {seg['start']:.1f}s">▶️</button>
                </div>
                """
                
            combined_html = f"""
            <!DOCTYPE html>
            <html>
            <head>
            <style>
                body {{
                    font-family: "Source Sans Pro", sans-serif;
                    color: var(--text-color, #fafafa);
                    background-color: transparent;
                    margin: 0;
                    padding: 0;
                }}
                .container {{
                    display: flex;
                    flex-direction: column;
                    height: 95vh;
                }}
                .audio-container {{
                    position: sticky;
                    top: 0;
                    background-color: var(--background-color, #0e1117);
                    padding: 10px 0;
                    z-index: 100;
                    border-bottom: 1px solid var(--secondary-background-color, #333);
                }}
                .transcript-container {{
                    flex: 1;
                    overflow-y: auto;
                    padding-top: 10px;
                    padding-bottom: 20px;
                }}
                .segment {{
                    display: flex;
                    padding: 10px;
                    margin-bottom: 5px;
                    border-radius: 8px;
                    cursor: pointer;
                    transition: background-color 0.2s;
                }}
                .segment:hover {{
                    background-color: var(--secondary-background-color, #262730);
                }}
                .segment.active {{
                    background-color: rgba(255, 193, 7, 0.2);
                    border-left: 4px solid #ffc107;
                }}
                .time {{
                    width: 60px;
                    font-weight: bold;
                    color: var(--text-color, #999);
                    opacity: 0.7;
                }}
                .text {{
                    flex: 1;
                    padding-right: 10px;
                }}
                .btn {{
                    background: transparent;
                    border: none;
                    cursor: pointer;
                    font-size: 1.2em;
                    opacity: 0.7;
                    padding: 0;
                }}
                .btn:hover {{
                    opacity: 1;
                }}
            </style>
            </head>
            <body>
                <div class="container">
                    <div class="audio-container">
                        <audio id="audio-player" controls style="width: 100%;">
                            <source src="data:{audio_mime};base64,{audio_base64}" type="{audio_mime}">
                            Tu navegador no soporta audio.
                        </audio>
                    </div>
                    
                    <div class="transcript-container" id="transcript">
                        {segments_html}
                    </div>
                </div>

                <script>
                    const audio = document.getElementById('audio-player');
                    const segmentsData = {segments_json};
                    
                    function jumpTo(time) {{
                        audio.currentTime = time;
                        audio.play();
                    }}
                    
                    audio.addEventListener('timeupdate', function() {{
                        const currentTime = audio.currentTime;
                        
                        for (let i = 0; i < segmentsData.length; i++) {{
                            const seg = segmentsData[i];
                            if (currentTime >= seg.start && currentTime <= seg.end) {{
                                const el = document.getElementById('segment_' + i);
                                if (el && !el.classList.contains('active')) {{
                                    document.querySelectorAll('.segment').forEach(s => s.classList.remove('active'));
                                    el.classList.add('active');
                                    el.scrollIntoView({{ behavior: 'smooth', block: 'center' }});
                                }}
                                break;
                            }}
                        }}
                    }});
                </script>
            </body>
            </html>
            """
            
            st.iframe(combined_html, height=600, scrolling=False)
            
            # Botón para limpiar
            if st.button("🗑️ Limpiar todo", use_container_width=True):
                if "audio_path" in st.session_state:
                    try:
                        os.unlink(st.session_state.audio_path)
                    except:
                        pass
                keys = ["segments", "detected_lang", "audio_path", "full_text", "last_answer", "last_question"]
                for key in keys:
                    if key in st.session_state:
                        del st.session_state[key]
                st.rerun()
        
        with col2:
            st.subheader("💬 Consultar a la IA")
            
            question = st.text_area(
                "Haz una pregunta sobre el audio:",
                placeholder="Ej: ¿Cuál es el tema principal? ¿Qué opinión dio el orador?",
                height=100
            )
            
            if st.button("🔍 Preguntar", type="primary", use_container_width=True):
                if question.strip():
                    with st.spinner("Consultando IA local..."):
                        model, tokenizer = load_llm()
                        answer = chat_with_context(
                            question,
                            st.session_state.segments,
                            model,
                            tokenizer
                        )
                        st.session_state.last_answer = answer
                        st.session_state.last_question = question
                else:
                    st.warning("Escribe una pregunta")
            
            if "last_answer" in st.session_state:
                st.markdown("---")
                st.markdown("**🤖 Respuesta:**")
                st.info(st.session_state.last_answer)
            
            # Mostrar estadísticas
            if "segments" in st.session_state:
                st.markdown("---")
                st.markdown("**📊 Estadísticas**")
                total_duration = st.session_state.segments[-1]["end"] if st.session_state.segments else 0
                st.metric("Duración total", f"{total_duration:.1f} segundos")
                st.metric("Segmentos", len(st.session_state.segments))
                st.metric("Idioma detectado", st.session_state.detected_lang.upper())
    
    else:
        st.info("👈 **Sube un archivo de audio en la barra lateral para comenzar**")

if __name__ == "__main__":
    main()