Update app.py

app.py

@@ -8,12 +8,29 @@ import librosa
 import soundfile as sf
 import pandas as pd
 import numpy as np
-from typing import List, Tuple
+from typing import List, Dict
 import shutil
+import threading
+import time
+from datetime import datetime
+import json
+import traceback
 
 # Khởi tạo model Whisper
 model = WhisperModel("large-v3-turbo", device="cpu", compute_type="int8")
 
+# Lưu trữ tasks và history
+TASKS = {}
+TASK_LOCK = threading.Lock()
+STORAGE_DIR = "task_storage"
+os.makedirs(STORAGE_DIR, exist_ok=True)
+
+class TaskStatus:
+    WAITING = "⏳ Đang chờ"
+    PROCESSING = "🔄 Đang xử lý"
+    SUCCESS = "✅ Thành công"
+    ERROR = "❌ Lỗi"
+
 def extract_audio_files(input_file: str, temp_dir: str) -> List[str]:
     """Giải nén file zip hoặc copy file audio đơn"""
     audio_files = []
@@ -64,7 +81,6 @@ def merge_short_segments(segments: List[dict], min_duration: float = 2.0) -> Lis
         current_duration = current['end'] - current['start']
         
         if current_duration < min_duration:
-            # Gộp với segment tiếp theo
             current['end'] = seg['end']
             current['text'] = current['text'] + ' ' + seg['text']
         else:
@@ -121,12 +137,10 @@ def save_to_parquet(records: List[dict], output_dir: str, max_size_mb: int = 500
     parquet_files = []
     
     if file_size_mb <= max_size_mb:
-        # Lưu thành 1 file
-        output_path = os.path.join(output_dir, 'dataset.parquet')
+        output_path = os.path.join(output_dir, 'train-00000-of-00001.parquet')
         df.to_parquet(output_path, engine='pyarrow')
         parquet_files.append(output_path)
     else:
-        # Chia nhỏ thành nhiều parts
         num_parts = int(np.ceil(file_size_mb / max_size_mb))
         chunk_size = len(df) // num_parts + 1
         
@@ -135,22 +149,38 @@ def save_to_parquet(records: List[dict], output_dir: str, max_size_mb: int = 500
             end_idx = min((i + 1) * chunk_size, len(df))
             
             df_chunk = df.iloc[start_idx:end_idx]
-            output_path = os.path.join(output_dir, f'dataset_part{i+1:03d}.parquet')
+            output_path = os.path.join(output_dir, f'train-{i:05d}-of-{num_parts:05d}.parquet')
             df_chunk.to_parquet(output_path, engine='pyarrow')
             parquet_files.append(output_path)
     
     return parquet_files
 
-def process_audio(input_file):
-    """Xử lý chính"""
-    if input_file is None:
-        return None, "Vui lòng upload file audio hoặc file zip!"
-    
-    with tempfile.TemporaryDirectory() as temp_dir:
-        # Tạo thư mục con
-        extract_dir = os.path.join(temp_dir, 'extracted')
-        audio_output_dir = os.path.join(temp_dir, 'audio')
-        final_output_dir = os.path.join(temp_dir, 'output')
+def update_task_status(task_id: str, status: str, details: dict = None):
+    """Cập nhật trạng thái task"""
+    with TASK_LOCK:
+        if task_id in TASKS:
+            TASKS[task_id]['status'] = status
+            TASKS[task_id]['updated_at'] = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            if details:
+                TASKS[task_id].update(details)
+            
+            # Lưu vào file
+            with open(os.path.join(STORAGE_DIR, f"{task_id}.json"), 'w', encoding='utf-8') as f:
+                json.dump(TASKS[task_id], f, ensure_ascii=False, indent=2)
+
+def process_audio_background(task_id: str, input_file: str, original_filename: str):
+    """Xử lý audio trong background"""
+    try:
+        update_task_status(task_id, TaskStatus.PROCESSING, {
+            'progress': 'Đang giải nén và phát hiện file audio...'
+        })
+        
+        task_dir = os.path.join(STORAGE_DIR, task_id)
+        os.makedirs(task_dir, exist_ok=True)
+        
+        extract_dir = os.path.join(task_dir, 'extracted')
+        audio_output_dir = os.path.join(task_dir, 'audio')
+        final_output_dir = os.path.join(task_dir, 'output')
         
         os.makedirs(extract_dir, exist_ok=True)
         os.makedirs(audio_output_dir, exist_ok=True)
@@ -160,12 +190,24 @@ def process_audio(input_file):
         audio_files = extract_audio_files(input_file, extract_dir)
         
         if not audio_files:
-            return None, "Không tìm thấy file audio nào!"
+            update_task_status(task_id, TaskStatus.ERROR, {
+                'error': 'Không tìm thấy file audio nào trong file tải lên!'
+            })
+            return
+        
+        update_task_status(task_id, TaskStatus.PROCESSING, {
+            'progress': f'Tìm thấy {len(audio_files)} file audio. Đang transcribe...',
+            'total_files': len(audio_files)
+        })
         
         all_records = []
         
         # Xử lý từng file audio
-        for audio_file in audio_files:
+        for idx, audio_file in enumerate(audio_files):
+            update_task_status(task_id, TaskStatus.PROCESSING, {
+                'progress': f'Đang xử lý file {idx+1}/{len(audio_files)}: {Path(audio_file).name}'
+            })
+            
             base_name = Path(audio_file).stem
             
             # Transcribe
@@ -184,6 +226,10 @@ def process_audio(input_file):
             
             all_records.extend(records)
         
+        update_task_status(task_id, TaskStatus.PROCESSING, {
+            'progress': f'Đã tạo {len(all_records)} segments. Đang lưu vào Parquet...'
+        })
+        
         # Lưu vào parquet
         parquet_files = save_to_parquet(all_records, final_output_dir)
         
@@ -192,7 +238,7 @@ def process_audio(input_file):
         shutil.copytree(audio_output_dir, final_audio_dir)
         
         # Tạo file zip
-        zip_path = os.path.join(temp_dir, 'dataset_output.zip')
+        zip_path = os.path.join(task_dir, 'dataset_output.zip')
         with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
             # Thêm audio files
             for root, _, files in os.walk(final_audio_dir):
@@ -205,57 +251,206 @@ def process_audio(input_file):
             for pq_file in parquet_files:
                 zipf.write(pq_file, os.path.basename(pq_file))
         
-        # Copy sang vị trí tạm để Gradio có thể trả về
-        final_zip = os.path.join(tempfile.gettempdir(), 'dataset_output.zip')
-        shutil.copy(zip_path, final_zip)
+        # Tính kích thước file
+        zip_size_mb = os.path.getsize(zip_path) / (1024 * 1024)
         
-        summary = f"""
-✅ Xử lý thành công!
-- Số file audio đầu vào: {len(audio_files)}
-- Số segment đã tạo: {len(all_records)}
-- Số file parquet: {len(parquet_files)}
-- File zip đầu ra: dataset_output.zip
-"""
+        update_task_status(task_id, TaskStatus.SUCCESS, {
+            'progress': 'Hoàn thành!',
+            'input_files': len(audio_files),
+            'total_segments': len(all_records),
+            'parquet_files': len(parquet_files),
+            'output_zip': zip_path,
+            'zip_size_mb': round(zip_size_mb, 2)
+        })
+        
+    except Exception as e:
+        error_msg = f"{str(e)}\n\n{traceback.format_exc()}"
+        update_task_status(task_id, TaskStatus.ERROR, {
+            'error': error_msg
+        })
+
+def submit_task(input_file):
+    """Submit task mới"""
+    if input_file is None:
+        return "❌ Vui lòng upload file audio hoặc file zip!", ""
+    
+    task_id = f"task_{int(time.time() * 1000)}"
+    original_filename = Path(input_file).name
+    
+    task_info = {
+        'task_id': task_id,
+        'status': TaskStatus.WAITING,
+        'created_at': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        'updated_at': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        'original_filename': original_filename,
+        'progress': 'Task đã được tạo, đang chờ xử lý...'
+    }
+    
+    with TASK_LOCK:
+        TASKS[task_id] = task_info
+    
+    # Chạy background thread
+    thread = threading.Thread(
+        target=process_audio_background, 
+        args=(task_id, input_file, original_filename),
+        daemon=True
+    )
+    thread.start()
+    
+    return f"✅ Task {task_id} đã được tạo và đang xử lý trong background!", task_id
+
+def load_all_tasks():
+    """Load tất cả tasks từ storage"""
+    with TASK_LOCK:
+        for file in os.listdir(STORAGE_DIR):
+            if file.endswith('.json'):
+                task_id = file.replace('.json', '')
+                if task_id not in TASKS:
+                    with open(os.path.join(STORAGE_DIR, file), 'r', encoding='utf-8') as f:
+                        TASKS[task_id] = json.load(f)
+
+def get_task_list():
+    """Lấy danh sách tasks để hiển thị trong dropdown"""
+    load_all_tasks()
+    with TASK_LOCK:
+        task_list = [(f"{task['task_id']} - {task['status']} - {task['original_filename']}", 
+                      task['task_id']) 
+                     for task in sorted(TASKS.values(), 
+                                       key=lambda x: x['created_at'], 
+                                       reverse=True)]
+    return task_list
+
+def get_task_info(task_id):
+    """Lấy thông tin chi tiết của task"""
+    if not task_id:
+        return "Chọn task để xem thông tin", None
+    
+    load_all_tasks()
+    
+    with TASK_LOCK:
+        if task_id not in TASKS:
+            return "Task không tồn tại!", None
         
-        return final_zip, summary
+        task = TASKS[task_id]
+    
+    info = f"""
+## 📋 Thông tin Task: {task_id}
+
+**Trạng thái:** {task['status']}
+**File gốc:** {task.get('original_filename', 'N/A')}
+**Thời gian tạo:** {task['created_at']}
+**Cập nhật lần cuối:** {task['updated_at']}
+
+---
+
+### 📊 Chi tiết
+
+**Tiến trình:** {task.get('progress', 'N/A')}
+"""
+    
+    if task['status'] == TaskStatus.SUCCESS:
+        info += f"""
+**Số file audio đầu vào:** {task.get('input_files', 'N/A')}
+**Tổng số segments:** {task.get('total_segments', 'N/A')}
+**Số file Parquet:** {task.get('parquet_files', 'N/A')}
+**Kích thước ZIP:** {task.get('zip_size_mb', 'N/A')} MB
+"""
+        zip_path = task.get('output_zip')
+        if zip_path and os.path.exists(zip_path):
+            return info, zip_path
+    
+    elif task['status'] == TaskStatus.ERROR:
+        info += f"""
+**Lỗi:**
+```
+{task.get('error', 'Unknown error')}
+```
+"""
+    
+    return info, None
+
+def refresh_task_list():
+    """Refresh danh sách tasks"""
+    choices = get_task_list()
+    return gr.Dropdown(choices=choices, value=choices[0][1] if choices else None)
+
+# Load tasks khi khởi động
+load_all_tasks()
 
 # Tạo giao diện Gradio
-with gr.Blocks(title="Audio Transcription & Dataset Creator") as app:
+with gr.Blocks(title="Audio Transcription & Dataset Creator", theme=gr.themes.Soft()) as app:
     gr.Markdown("""
-    # 🎙️ Audio Transcription & Dataset Creator
-    Upload file audio hoặc file zip chứa nhiều file audio. 
-    Hệ thống sẽ:
-    1. Transcribe bằng Whisper Large-v3-Turbo
-    2. Cắt audio theo timestamps (gộp câu ngắn)
-    3. Tạo dataset Parquet chuẩn với audio bytes
+    # 🎙️ Audio Transcription & Dataset Creator with Background Processing
+    Upload file audio hoặc file zip - Hệ thống xử lý trong background và lưu lịch sử
     """)
     
-    with gr.Row():
-        with gr.Column():
-            input_file = gr.File(
-                label="Upload Audio File hoặc ZIP",
-                file_types=['.wav', '.mp3', '.flac', '.ogg', '.m4a', '.aac', '.zip']
-            )
-            process_btn = gr.Button("🚀 Bắt đầu xử lý", variant="primary")
+    with gr.Tabs():
+        # Tab Upload
+        with gr.Tab("📤 Upload & Submit"):
+            gr.Markdown("### Tải lên file và submit task")
+            
+            with gr.Row():
+                with gr.Column():
+                    input_file = gr.File(
+                        label="Upload Audio File hoặc ZIP",
+                        file_types=['.wav', '.mp3', '.flac', '.ogg', '.m4a', '.aac', '.zip']
+                    )
+                    submit_btn = gr.Button("🚀 Submit Task", variant="primary", size="lg")
+                
+                with gr.Column():
+                    submit_status = gr.Textbox(label="📋 Trạng thái Submit", lines=3)
+                    current_task_id = gr.Textbox(label="Task ID", visible=False)
+            
+            gr.Markdown("""
+            ### ℹ️ Hướng dẫn:
+            1. Upload file audio hoặc ZIP chứa nhiều file audio
+            2. Click "Submit Task" - task sẽ chạy trong background
+            3. Chuyển sang tab "History" để xem tiến trình và tải kết quả
+            """)
         
-        with gr.Column():
-            output_file = gr.File(label="📦 Tải về Dataset ZIP")
-            status_text = gr.Textbox(label="📊 Trạng thái", lines=8)
+        # Tab History
+        with gr.Tab("📜 History"):
+            gr.Markdown("### Xem lại lịch sử tasks và tải kết quả")
+            
+            with gr.Row():
+                refresh_btn = gr.Button("🔄 Refresh", size="sm")
+                task_dropdown = gr.Dropdown(
+                    label="Chọn Task",
+                    choices=get_task_list(),
+                    value=get_task_list()[0][1] if get_task_list() else None,
+                    interactive=True
+                )
+            
+            task_info_display = gr.Markdown("Chọn task để xem thông tin")
+            
+            download_btn = gr.File(label="📦 Tải về Dataset ZIP")
+            
+            # Auto refresh mỗi 3 giây
+            timer = gr.Timer(3)
     
-    process_btn.click(
-        fn=process_audio,
+    # Event handlers
+    submit_btn.click(
+        fn=submit_task,
         inputs=input_file,
-        outputs=[output_file, status_text]
+        outputs=[submit_status, current_task_id]
     )
     
-    gr.Markdown("""
-    ### 📝 Ghi chú:
-    - Dataset Parquet sẽ được chia nhỏ nếu > 500MB
-    - Cột `audio`: audio bytes (binary)
-    - Cột `transcription`: văn bản transcription
-    - Cột `file_name`: đường dẫn dạng `audio/filename_00001.wav`
-    - Các câu ngắn (< 2s) sẽ được gộp lại
-    """)
+    refresh_btn.click(
+        fn=refresh_task_list,
+        outputs=task_dropdown
+    )
+    
+    task_dropdown.change(
+        fn=get_task_info,
+        inputs=task_dropdown,
+        outputs=[task_info_display, download_btn]
+    )
+    
+    timer.tick(
+        fn=get_task_info,
+        inputs=task_dropdown,
+        outputs=[task_info_display, download_btn]
+    )
 
 if __name__ == "__main__":
     app.launch()