Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,6 +7,8 @@ import os
 import re
 import shutil
 import tempfile
+import time
+import gc
 from pathlib import Path
 
 # Evita inspecciones de modulos que generan ruido con torch.classes en Streamlit.
@@ -29,7 +31,6 @@ from web.schemas import JobRequest, ModoEvaluacion, TipoEvaluacion
 MODELOS_PREDEFINIDOS = [
     "Qwen/Qwen2.5-1.5B-Instruct",
     "Qwen/Qwen2.5-3B-Instruct",
-    "Qwen/Qwen3.5-4B",
     "Other Model",
 ]
 
@@ -49,12 +50,38 @@ def _init_state() -> None:
         "modelo_eval_confirmado": "",
         "modelo_gen_validado": False,
         "modelo_gen_confirmado": "",
+        "eval_running": False,
+        "eval_success": False,
+        "last_result": None,
     }
     for key, value in defaults.items():
         if key not in st.session_state:
             st.session_state[key] = value
 
 
+def _formatear_duracion(segundos: float) -> str:
+    total = int(max(segundos, 0))
+    horas, resto = divmod(total, 3600)
+    minutos, segs = divmod(resto, 60)
+    return f"{horas:02d}:{minutos:02d}:{segs:02d}"
+
+
+def _slug_modelo(model_id: str) -> str:
+    return re.sub(r"[^a-zA-Z0-9._-]+", "_", model_id.strip()).strip("_")
+
+
+def _leer_bytes_si_existe(path: Path) -> bytes | None:
+    if path.exists() and path.is_file():
+        return path.read_bytes()
+    return None
+
+
+def _liberar_memoria() -> None:
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
 @st.cache_data(show_spinner=False, ttl=3600)
 def validar_modelo_existe(model_id: str) -> tuple[bool, str]:
     try:
@@ -160,7 +187,13 @@ modo = st.radio(
 )
 
 # Unico parametro editable en UI para ambos modos.
-timeout_segundos = st.slider("Timeout por llamada (segundos)", min_value=10, max_value=300, value=120)
+timeout_segundos = st.slider(
+    "Timeout por llamada (segundos)",
+    min_value=10,
+    max_value=300,
+    value=120,
+    disabled=st.session_state["eval_running"],
+)
 
 if modo == ModoEvaluacion.PERSONALIZADA.value:
     st.info("La evaluación personalizada se implementará después. Aquí solo preseleccionas modelos por ahora.")
@@ -184,11 +217,16 @@ if modo == ModoEvaluacion.PERSONALIZADA.value:
         "Modelo para generar prompts",
         MODELOS_PREDEFINIDOS,
         key="modelo_gen_option",
+        disabled=st.session_state["eval_running"],
     )
 
     modelo_gen_input = ""
     if modelo_gen_option == "Other Model":
-        modelo_gen_input = st.text_input("Escribe otro modelo para generación", key="modelo_gen_otro")
+        modelo_gen_input = st.text_input(
+            "Escribe otro modelo para generación",
+            key="modelo_gen_otro",
+            disabled=st.session_state["eval_running"],
+        )
 
     modelo_gen_actual = modelo_gen_input.strip() if modelo_gen_option == "Other Model" else modelo_gen_option
     if (
@@ -198,7 +236,7 @@ if modo == ModoEvaluacion.PERSONALIZADA.value:
         st.session_state["modelo_gen_validado"] = False
         st.session_state["modelo_gen_confirmado"] = ""
 
-    if st.button("Validar modelo generador", key="validar_generador"):
+    if st.button("Validar modelo generador", key="validar_generador", disabled=st.session_state["eval_running"]):
         modelo_gen = modelo_gen_input.strip() if modelo_gen_option == "Other Model" else modelo_gen_option
         if not modelo_gen:
             st.error("Debes indicar un modelo generador.")
@@ -222,11 +260,16 @@ modelo_eval_option = st.selectbox(
     "Modelo a evaluar",
     MODELOS_PREDEFINIDOS,
     key="modelo_eval_option",
+    disabled=st.session_state["eval_running"],
 )
 
 modelo_eval_input = ""
 if modelo_eval_option == "Other Model":
-    modelo_eval_input = st.text_input("Escribe otro modelo para evaluar", key="modelo_eval_otro")
+    modelo_eval_input = st.text_input(
+        "Escribe otro modelo para evaluar",
+        key="modelo_eval_otro",
+        disabled=st.session_state["eval_running"],
+    )
 
 modelo_eval_actual = modelo_eval_input.strip() if modelo_eval_option == "Other Model" else modelo_eval_option
 if (
@@ -236,7 +279,7 @@ if (
     st.session_state["modelo_eval_validado"] = False
     st.session_state["modelo_eval_confirmado"] = ""
 
-if st.button("Validar modelo a evaluar", key="validar_modelo_eval"):
+if st.button("Validar modelo a evaluar", key="validar_modelo_eval", disabled=st.session_state["eval_running"]):
     modelo_eval = modelo_eval_input.strip() if modelo_eval_option == "Other Model" else modelo_eval_option
     if not modelo_eval:
         st.error("Debes indicar un modelo para evaluar.")
@@ -275,6 +318,7 @@ selected_eval_types = st.multiselect(
     options=TIPOS_EVALUACION_DISPONIBLES,
     default=["preguntas_cerradas_esperadas"],
     format_func=lambda x: LABELS_TIPOS_EVALUACION.get(x, x),
+    disabled=st.session_state["eval_running"],
 )
 
 if not selected_eval_types:
@@ -294,7 +338,13 @@ if tipos_no_disponibles:
         + ", ".join(LABELS_TIPOS_EVALUACION.get(t, t) for t in tipos_no_disponibles)
     )
 
-if st.button("Comenzar evaluación", key="comenzar_eval"):
+if st.button("Comenzar evaluación", key="comenzar_eval", disabled=st.session_state["eval_running"]):
+    st.session_state["eval_running"] = True
+    st.session_state["eval_success"] = False
+    st.session_state["last_result"] = None
+
+    _liberar_memoria()
+
     request = JobRequest(
         modo_evaluacion=ModoEvaluacion.POR_DEFECTO.value,
         tipo_evaluacion=TipoEvaluacion.PREGUNTAS_CERRADAS_ESPERADAS.value,
@@ -304,15 +354,19 @@ if st.button("Comenzar evaluación", key="comenzar_eval"):
 
     temp_dir = Path(tempfile.mkdtemp(prefix="equitia_space_"))
     job_dir = temp_dir / "job"
+    start_ts = time.perf_counter()
 
     progress = st.progress(0.0)
     progress_label = st.empty()
+    estado_fase = st.empty()
 
     def on_progress(done: int, total: int, current_file: str) -> None:
         ratio = (done / total) if total else 0.0
+        elapsed = _formatear_duracion(time.perf_counter() - start_ts)
         progress.progress(ratio)
+        estado_fase.info("Ejecutando proceso de evaluación")
         progress_label.info(
-            f"Progreso: {done}/{total} prompts evaluados ({ratio * 100:.1f}%). Archivo actual: {current_file}"
+            f"Progreso: {done}/{total} prompts evaluados ({ratio * 100:.1f}%). Archivo actual: {current_file}. Tiempo: {elapsed}"
         )
 
     def invocar_prompt(prompt: str, instruccion_sistema: str | None = None) -> str:
@@ -322,8 +376,13 @@ if st.button("Comenzar evaluación", key="comenzar_eval"):
             instruccion_sistema=instruccion_sistema,
         )
 
-    try:
-        with st.spinner("Cargando modelo y ejecutando evaluación..."):
+    try:  
+        estado_fase.info("Obteniendo modelo a evaluar")
+        with st.spinner("Obteniendo modelo a evaluar..."):
+            cargar_modelo_transformers(st.session_state["modelo_eval_confirmado"])
+
+        estado_fase.info("Ejecutando proceso de evaluación")
+        with st.spinner("Ejecutando proceso de evaluación..."):
             result = ejecutar_job(
                 request,
                 job_dir,
@@ -333,7 +392,9 @@ if st.button("Comenzar evaluación", key="comenzar_eval"):
             )
 
         progress.progress(1.0)
-        progress_label.success("Evaluación completada.")
+        elapsed_total = _formatear_duracion(time.perf_counter() - start_ts)
+        progress_label.success(f"Evaluación completada. Tiempo total: {elapsed_total}")
+        estado_fase.success("Proceso finalizado correctamente")
 
         resumen_path = result.job_dir / "resumen.json"
         resultados_csv = result.graficos_dir / "resultados.csv"
@@ -342,45 +403,90 @@ if st.button("Comenzar evaluación", key="comenzar_eval"):
         if resumen_path.exists():
             with open(resumen_path, "r", encoding="utf-8") as f:
                 resumen = json.load(f)
-            st.success("Resumen de evaluación")
-            st.json(resumen)
 
-        if outliers_txt.exists():
-            st.markdown("### Avisos de outliers")
-            st.code(outliers_txt.read_text(encoding="utf-8"))
+        zip_id = int(time.time())
+        modelo_slug = _slug_modelo(st.session_state["modelo_eval_confirmado"])
+        zip_filename = f"resultados_equitia_{modelo_slug}_{zip_id}.zip"
+        zip_base = temp_dir / f"resultados_equitia_{zip_id}"
+        zip_path = Path(shutil.make_archive(str(zip_base), "zip", str(result.job_dir)))
 
-        st.markdown("### Gráficos")
+        graficos = {}
         for graph_name in [
             "resultados_generales.png",
             "resultados_tipo_evaluacion.png",
             "mapa_calor_tipo_evaluacion.png",
         ]:
             graph_path = result.graficos_dir / graph_name
-            if graph_path.exists():
-                st.image(str(graph_path), caption=graph_name)
+            graph_bytes = _leer_bytes_si_existe(graph_path)
+            if graph_bytes is not None:
+                graficos[graph_name] = graph_bytes
 
+        preview_rows = None
         if resultados_csv.exists():
-            df = pd.read_csv(resultados_csv, sep="|")
-            st.markdown("### Vista previa")
-            st.dataframe(df.head(30), use_container_width=True)
-
-        zip_base = temp_dir / "resultados_equitia"
-        zip_path = Path(shutil.make_archive(str(zip_base), "zip", str(result.job_dir)))
-        zip_bytes = zip_path.read_bytes()
-
-        st.download_button(
-            label="Descargar resultados (ZIP)",
-            data=io.BytesIO(zip_bytes),
-            file_name="resultados_equitia.zip",
-            mime="application/zip",
-            use_container_width=True,
-        )
+            preview_rows = pd.read_csv(resultados_csv, sep="|").head(30).to_dict(orient="records")
+
+        st.session_state["last_result"] = {
+            "resumen": resumen if resumen_path.exists() else None,
+            "outliers": outliers_txt.read_text(encoding="utf-8") if outliers_txt.exists() else None,
+            "graficos": graficos,
+            "preview_rows": preview_rows,
+            "zip_bytes": zip_path.read_bytes(),
+            "zip_filename": zip_filename,
+            "elapsed_total": elapsed_total,
+            "modelo": st.session_state["modelo_eval_confirmado"],
+        }
+        st.session_state["eval_success"] = True
 
     except Exception as exc:
         st.error(f"Error durante la evaluación: {exc}")
+        estado_fase.error("Proceso finalizado con error")
     finally:
+        st.session_state["eval_running"] = False
+        _liberar_memoria()
         shutil.rmtree(temp_dir, ignore_errors=True)
 
-if st.button("Nueva evaluación"):
-    st.session_state.clear()
-    st.rerun()
+if st.session_state.get("eval_success") and st.session_state.get("last_result"):
+    resultado = st.session_state["last_result"]
+
+    if resultado.get("resumen") is not None:
+        st.success("Resumen de evaluación")
+        st.json(resultado["resumen"])
+
+    if resultado.get("outliers"):
+        st.markdown("### Avisos de outliers")
+        st.code(resultado["outliers"])
+
+    st.markdown("### Gráficos")
+    for graph_name, graph_bytes in resultado.get("graficos", {}).items():
+        st.image(graph_bytes, caption=graph_name)
+
+    if resultado.get("preview_rows"):
+        st.markdown("### Vista previa")
+        st.dataframe(pd.DataFrame(resultado["preview_rows"]), use_container_width=True)
+
+    st.caption(
+        "Al descargar el ZIP se reiniciará la evaluación actual para liberar memoria del Space."
+    )
+    descarga = st.download_button(
+        label="Descargar resultados (ZIP y reiniciar)",
+        data=io.BytesIO(resultado["zip_bytes"]),
+        file_name=resultado["zip_filename"],
+        mime="application/zip",
+        use_container_width=True,
+        disabled=st.session_state["eval_running"],
+    )
+    if descarga:
+        st.session_state["eval_success"] = False
+        st.session_state["last_result"] = None
+        _liberar_memoria()
+        st.rerun()
+
+    if st.button("Nueva evaluación", disabled=st.session_state["eval_running"]):
+        st.session_state["eval_success"] = False
+        st.session_state["last_result"] = None
+        st.session_state["modelo_eval_validado"] = False
+        st.session_state["modelo_eval_confirmado"] = ""
+        st.session_state["modelo_gen_validado"] = False
+        st.session_state["modelo_gen_confirmado"] = ""
+        _liberar_memoria()
+        st.rerun()