Update app.py

app.py

@@ -79,19 +79,25 @@ server_lock = asyncio.Lock()  # Lock for thread-safe server state access
 
 def load_progress() -> Dict:
     """Loads the local processing progress from the JSON file."""
+    default_structure = {
+        "last_processed_index": 0,
+        "processed_files": {}, # {index: repo_path}
+        "file_list": [], # Full list of all zip files found in the dataset
+        "uploaded_count": 0
+    }
     if PROGRESS_FILE.exists():
         try:
             with PROGRESS_FILE.open('r') as f:
-                return json.load(f)
+                data = json.load(f)
+                # Ensure all keys exist
+                for key, value in default_structure.items():
+                    if key not in data:
+                        data[key] = value
+                return data
         except json.JSONDecodeError:
             print(f"[{FLOW_ID}] WARNING: Progress file is corrupted. Starting fresh.")
     
-    return {
-        "last_processed_index": 0,
-        "processed_files": {}, # {index: repo_path}
-        "file_list": [], # Full list of all zip files found in the dataset
-        "uploaded_count": 0
-    }
+    return default_structure
 
 def save_progress(progress_data: Dict):
     """Saves the local processing progress to the JSON file."""
@@ -159,7 +165,7 @@ async def upload_hf_state(state: Dict[str, Any]) -> bool:
 # --- Hugging Face Utility Functions ---
 
 async def get_audio_file_list(progress_data: Dict) -> List[str]:
-    if progress_data['file_list']:
+    if progress_data.get('file_list'):
         return progress_data['file_list']
     try:
         api = HfApi(token=HF_TOKEN)
@@ -228,7 +234,6 @@ async def process_file_task(wav_file: str, state: Dict, progress: Dict):
             state["file_states"][wav_file] = "processed"
             progress["uploaded_count"] = progress.get("uploaded_count", 0) + 1
             print(f"[{FLOW_ID}] ✅ Success: {wav_file}")
-            # Note: In a real scenario, you'd save the 'result' (transcription) somewhere
         else:
             state["file_states"][wav_file] = "failed_transcription"
             print(f"[{FLOW_ID}] ❌ Failed: {wav_file}")
@@ -246,51 +251,54 @@ async def main_processing_loop():
     print(f"[{FLOW_ID}] Starting main processing loop...")
     
     while True:
-        state = await download_hf_state()
-        progress = load_progress()
-        file_list = await get_audio_file_list(progress)
-        
-        if not file_list:
-            await asyncio.sleep(60)
-            continue
-
-        # 1. Handpick failed_transcription files
-        failed_files = [f for f, s in state.get("file_states", {}).items() if s == "failed_transcription"]
-        
-        # 2. Also check for new files based on next_download_index
-        next_idx = state.get("next_download_index", 0)
-        new_files = file_list[next_idx:next_idx + 100] # Take a chunk of new files
-        
-        # Combine: Prioritize failed files, then add new ones
-        files_to_process = failed_files + [f for f in new_files if f not in state["file_states"]]
-        
-        if not files_to_process:
-            print(f"[{FLOW_ID}] No files to process. Sleeping...")
-            await asyncio.sleep(60)
-            continue
+        try:
+            state = await download_hf_state()
+            progress = load_progress()
+            file_list = await get_audio_file_list(progress)
+            
+            if not file_list:
+                print(f"[{FLOW_ID}] File list empty, retrying in 60s...")
+                await asyncio.sleep(60)
+                continue
 
-        print(f"[{FLOW_ID}] Processing {len(files_to_process)} files ({len(failed_files)} failed, {len(files_to_process)-len(failed_files)} new)...")
-        
-        # Process in batches of server count
-        batch_size = len(servers)
-        for i in range(0, len(files_to_process), batch_size):
-            batch = files_to_process[i:i + batch_size]
-            tasks = [process_file_task(f, state, progress) for f in batch]
-            await asyncio.gather(*tasks)
+            # 1. Handpick failed_transcription files
+            failed_files = [f for f, s in state.get("file_states", {}).items() if s == "failed_transcription"]
             
-            # Update next_download_index if we processed new files
-            processed_new = [f for f in batch if f in new_files]
-            if processed_new:
-                # This is a simple way to update index; in reality, you'd want to be more precise
-                # but for this fix, we'll just increment based on what we found
-                last_new_file = processed_new[-1]
-                state["next_download_index"] = file_list.index(last_new_file) + 1
+            # 2. Also check for new files based on next_download_index
+            next_idx = state.get("next_download_index", 0)
+            new_files = file_list[next_idx:next_idx + 100] # Take a chunk of new files
             
-            # Save and upload state after each batch
-            await upload_hf_state(state)
-            save_progress(progress)
+            # Combine: Prioritize failed files, then add new ones
+            files_to_process = failed_files + [f for f in new_files if f not in state["file_states"]]
             
-        await asyncio.sleep(10)
+            if not files_to_process:
+                print(f"[{FLOW_ID}] No files to process. Sleeping...")
+                await asyncio.sleep(60)
+                continue
+
+            print(f"[{FLOW_ID}] Processing {len(files_to_process)} files ({len(failed_files)} failed, {len(files_to_process)-len(failed_files)} new)...")
+            
+            # Process in batches of server count
+            batch_size = len(servers)
+            for i in range(0, len(files_to_process), batch_size):
+                batch = files_to_process[i:i + batch_size]
+                tasks = [process_file_task(f, state, progress) for f in batch]
+                await asyncio.gather(*tasks)
+                
+                # Update next_download_index if we processed new files
+                processed_new = [f for f in batch if f in new_files]
+                if processed_new:
+                    last_new_file = processed_new[-1]
+                    state["next_download_index"] = file_list.index(last_new_file) + 1
+                
+                # Save and upload state after each batch
+                await upload_hf_state(state)
+                save_progress(progress)
+                
+            await asyncio.sleep(10)
+        except Exception as e:
+            print(f"[{FLOW_ID}] Error in main loop: {e}")
+            await asyncio.sleep(60)
 
 # --- FastAPI App ---
 
@@ -310,12 +318,12 @@ async def root():
         "status": "running",
         "next_download_index": state.get("next_download_index", 0),
         "failed_transcriptions": failed_count,
-        "uploaded_count": progress.get("uploaded_count", 0)
+        "uploaded_count": progress.get("uploaded_count", 0),
+        "total_files_in_list": len(progress.get('file_list', []))
     }
 
 @app.post("/start_processing")
 async def start_processing(request: ProcessStartRequest):
-    # This endpoint can be used to manually reset the index if needed
     state = await download_hf_state()
     state["next_download_index"] = request.start_index - 1
     await upload_hf_state(state)