Subida de código

whisper-distributed/unifier/app.py

@@ -0,0 +1,176 @@
+#!/usr/bin/env python
+import os
+import time
+import json
+import pika
+import logging
+from collections import defaultdict
+from datetime import datetime
+
+# Configuración de logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger('unifier')
+
+# Configuración de RabbitMQ
+RABBITMQ_HOST = 'rabbitmq'
+RABBITMQ_USER = 'user'
+RABBITMQ_PASS = 'password'
+UNIFY_QUEUE = 'text_unify_queue'
+
+# Directorios
+SHARED_DIR = '/app/shared'
+OUTPUT_DIR = '/app/output'
+
+# Almacenamiento temporal de transcripciones por trabajo
+pending_transcriptions = defaultdict(dict)
+job_metadata = {}
+
+def connect_to_rabbitmq():
+    """Establece conexión con RabbitMQ"""
+    tries = 0
+    while True:
+        try:
+            credentials = pika.PlainCredentials(RABBITMQ_USER, RABBITMQ_PASS)
+            connection = pika.BlockingConnection(
+                pika.ConnectionParameters(
+                    host=RABBITMQ_HOST,
+                    credentials=credentials
+                )
+            )
+            return connection
+        except pika.exceptions.AMQPConnectionError:
+            tries += 1
+            logger.warning(f"Intento {tries}: No se pudo conectar a RabbitMQ. Reintentando en 5 segundos...")
+            time.sleep(5)
+
+def process_transcription(segment_info):
+    """Procesa una transcripción recibida"""
+    original_file_id = segment_info['original_file_id']
+    segment_id = segment_info['segment_id']
+    transcription_path = segment_info['transcription_path']
+    total_segments = segment_info['total_segments']
+    
+    # Almacenar metadatos del trabajo si no existen
+    if original_file_id not in job_metadata:
+        job_metadata[original_file_id] = {
+            'original_filename': segment_info['original_filename'],
+            'total_segments': total_segments,
+            'start_time': datetime.now()
+        }
+    
+    # Leer transcripción
+    with open(transcription_path, 'r', encoding='utf-8') as f:
+        text = f.read()
+    
+    # Almacenar transcripción
+    pending_transcriptions[original_file_id][segment_id] = {
+        'text': text,
+        'start_ms': segment_info['start_ms'],
+        'end_ms': segment_info['end_ms']
+    }
+    
+    logger.info(f"Recibida transcripción {segment_id+1}/{total_segments} para {original_file_id}")
+    
+    # Verificar si se han recibido todas las transcripciones para este trabajo
+    if len(pending_transcriptions[original_file_id]) == total_segments:
+        unify_transcriptions(original_file_id)
+
+def convert_milliseconds_to_srt_time(ms):
+    """Convierte milisegundos a formato de tiempo SRT (HH:MM:SS,mmm)"""
+    seconds, milliseconds = divmod(ms, 1000)
+    minutes, seconds = divmod(seconds, 60)
+    hours, minutes = divmod(minutes, 60)
+    return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}"
+
+def unify_transcriptions(original_file_id):
+    """Unifica las transcripciones de un trabajo y genera el archivo de subtítulos"""
+    logger.info(f"Unificando transcripciones para {original_file_id}")
+    
+    # Obtener metadatos del trabajo
+    metadata = job_metadata[original_file_id]
+    original_filename = metadata['original_filename']
+    
+    # Preparar nombre de archivo de salida
+    base_filename = os.path.splitext(original_filename)[0]
+    output_filename = f"{base_filename}_transcription.srt"
+    output_path = os.path.join(OUTPUT_DIR, output_filename)
+    
+    # Ordenar transcripciones por segment_id
+    sorted_segments = sorted(pending_transcriptions[original_file_id].items(), key=lambda x: x[0])
+    
+    # Crear directorio de salida si no existe
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    
+    # Generar archivo SRT (formato de subtítulos)
+    with open(output_path, 'w', encoding='utf-8') as f:
+        for i, (segment_id, segment_data) in enumerate(sorted_segments):
+            # Formatear tiempo de inicio y fin
+            start_time = convert_milliseconds_to_srt_time(segment_data['start_ms'])
+            end_time = convert_milliseconds_to_srt_time(segment_data['end_ms'])
+            
+            # Escribir entrada de subtítulo
+            f.write(f"{i+1}\n")
+            f.write(f"{start_time} --> {end_time}\n")
+            f.write(f"{segment_data['text'].strip()}\n\n")
+    
+    # Calcular tiempo total de procesamiento
+    total_time = datetime.now() - metadata['start_time']
+    
+    logger.info(f"Archivo de subtítulos generado: {output_path}")
+    logger.info(f"Tiempo total de procesamiento: {total_time}")
+    
+    # Limpiar datos del trabajo completado
+    del pending_transcriptions[original_file_id]
+    del job_metadata[original_file_id]
+
+def callback(ch, method, properties, body):
+    """Callback para procesar mensajes de la cola de unificación"""
+    try:
+        # Decodificar mensaje
+        segment_info = json.loads(body)
+        
+        # Procesar transcripción
+        process_transcription(segment_info)
+        
+        # Confirmar procesamiento
+        ch.basic_ack(delivery_tag=method.delivery_tag)
+        
+    except Exception as e:
+        logger.error(f"Error procesando mensaje: {str(e)}")
+        # Rechazar mensaje en caso de error para reintentarlo
+        ch.basic_nack(delivery_tag=method.delivery_tag, requeue=True)
+
+def main():
+    """Función principal"""
+    # Asegurar que directorios existen
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    
+    # Conexión a RabbitMQ
+    connection = connect_to_rabbitmq()
+    channel = connection.channel()
+    
+    # Declarar cola
+    channel.queue_declare(queue=UNIFY_QUEUE, durable=True)
+    
+    # Configurar prefetch count
+    channel.basic_qos(prefetch_count=1)
+    
+    # Configurar callback
+    channel.basic_consume(queue=UNIFY_QUEUE, on_message_callback=callback)
+    
+    logger.info("Servicio unificador iniciado. Esperando transcripciones...")
+    
+    # Iniciar consumo
+    try:
+        channel.start_consuming()
+    except KeyboardInterrupt:
+        logger.info("Servicio unificador detenido")
+        channel.stop_consuming()
+    
+    connection.close()
+
+if __name__ == "__main__":
+    main()