retrieve

app.py

@@ -16,6 +16,55 @@ MODEL_FILE_PATH = os.path.join(MODEL_OUTPUT_DIR, MODEL_FILE_NAME)
 
 # ----------------------------------------------------------------
 
+def retrieve_model():
+    """
+    Checks for the final model file and prepares it for download.
+    Useful for when the training job finishes server-side but the
+    client connection has timed out.
+    """
+    MODEL_OUTPUT_DIR = "checkpoints"
+    MODEL_FILE_NAME = "layoutlmv3_crf_passage.pth"
+    MODEL_FILE_PATH = os.path.join(MODEL_OUTPUT_DIR, MODEL_FILE_NAME)
+
+    if os.path.exists(MODEL_FILE_PATH):
+        file_size = os.path.getsize(MODEL_FILE_PATH) / (1024 * 1024)  # Size in MB
+
+        # CRITICAL: Copy to a simple location that Gradio can reliably serve
+        import tempfile
+        temp_dir = tempfile.gettempdir()
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        temp_model_path = os.path.join(temp_dir, f"layoutlmv3_trained_{timestamp}_recovered.pth")
+
+        try:
+            shutil.copy2(MODEL_FILE_PATH, temp_model_path)
+            download_path = temp_model_path
+
+            log_output = (
+                f"--- Model Status Check: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ---\n"
+                f"🎉 SUCCESS! A trained model was found and recovered.\n"
+                f"📦 Model file: {MODEL_FILE_PATH}\n"
+                f"📊 Model size: {file_size:.2f} MB\n"
+                f"🔗 Download path prepared: {download_path}\n\n"
+                f"⬇️ Click the '📥 Download Model' button below to save your model."
+            )
+            return log_output, download_path, gr.Button(visible=True)
+
+        except Exception as e:
+            log_output = (
+                f"--- Model Status Check FAILED ---\n"
+                f"⚠️ Trained model found, but could not prepare for download: {e}\n"
+                f"📁 Original Path: {MODEL_FILE_PATH}. Try again or check Space logs."
+            )
+            return log_output, None, gr.Button(visible=False)
+
+    else:
+        log_output = (
+            f"--- Model Status Check: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ---\n"
+            f"❌ Model file not found at {MODEL_FILE_PATH}.\n"
+            f"Training may still be running or it failed. Check back later."
+        )
+        return log_output, None, gr.Button(visible=False)
+
 
 def train_model(dataset_file: gr.File, batch_size: int, epochs: int, lr: float, max_len: int, progress=gr.Progress()):
     """
@@ -256,6 +305,12 @@ with gr.Blocks(title="LayoutLMv3 Fine-Tuning App", theme=gr.themes.Soft()) as de
                 visible=False
             )
 
+            check_button.click(
+                fn=retrieve_model,  # A new function we'll define
+                inputs=[],
+                outputs=[log_output, model_path_state, download_btn]
+            )
+
             # File output for download
             model_download = gr.File(
                 label="Your trained model will appear here after clicking Download",