download .pth file

app.py

@@ -310,12 +310,173 @@
 #     # and sometimes cause issues in managed Space environments.
 #     demo.launch()
 
+#
+# import gradio as gr
+# import subprocess
+# import os
+# import sys
+# from datetime import datetime
+#
+# # FIX: Update the script name to the correct one you uploaded
+# TRAINING_SCRIPT = "HF_LayoutLM_with_Passage.py"
+#
+# # --- CORRECTED MODEL PATH BASED ON YOUR SCRIPT ---
+# MODEL_OUTPUT_DIR = "checkpoints"
+# MODEL_FILE_NAME = "layoutlmv3_crf_passage.pth"
+# MODEL_FILE_PATH = os.path.join(MODEL_OUTPUT_DIR, MODEL_FILE_NAME)
+#
+#
+# # ----------------------------------------------------------------
+#
+#
+# def train_model(dataset_file: gr.File, batch_size: int, epochs: int, lr: float, max_len: int, progress=gr.Progress()):
+#     """
+#     Handles the Gradio submission and executes the training script using subprocess.
+#     """
+#
+#     # 1. Setup: Create output directory if it doesn't exist
+#     os.makedirs(MODEL_OUTPUT_DIR, exist_ok=True)
+#
+#     # 2. File Handling: Use the temporary path of the uploaded file
+#     if dataset_file is None:
+#         yield "❌ ERROR: Please upload a file.", None
+#         return
+#
+#     # Using .name (Corrected in previous steps)
+#     input_path = dataset_file.name
+#
+#     if not input_path.lower().endswith(".json"):
+#         yield "❌ ERROR: Please upload a valid Label Studio JSON file (.json).", None
+#         return
+#
+#     progress(0.1, desc="Starting LayoutLMv3 Training...")
+#
+#     log_output = f"--- Training Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ---\n"
+#
+#     # 3. Construct the subprocess command
+#     command = [
+#         sys.executable,
+#         # Now uses the corrected TRAINING_SCRIPT variable
+#         TRAINING_SCRIPT,
+#         "--mode", "train",
+#         "--input", input_path,
+#         "--batch_size", str(batch_size),
+#         "--epochs", str(epochs),
+#         "--lr", str(lr),
+#         "--max_len", str(max_len)
+#     ]
+#
+#     log_output += f"Executing command: {' '.join(command)}\n\n"
+#     yield log_output, None  # Yield the command to the log output
+#
+#     try:
+#         # 4. Run the training script and capture output
+#         process = subprocess.Popen(
+#             command,
+#             stdout=subprocess.PIPE,
+#             stderr=subprocess.STDOUT,
+#             text=True,
+#             bufsize=1
+#         )
+#
+#         # Stream logs in real-time
+#         for line in iter(process.stdout.readline, ""):
+#             log_output += line
+#             yield log_output, None  # Send partial log to Gradio output
+#
+#         process.stdout.close()
+#         return_code = process.wait()
+#
+#         # 5. Check for successful completion
+#         if return_code == 0:
+#             log_output += "\n✅ TRAINING COMPLETE! Model saved."
+#
+#             # 6. Prepare download links based on script's saved path
+#             model_exists = os.path.exists(MODEL_FILE_PATH)
+#
+#             if model_exists:
+#                 log_output += f"\nModel path: {MODEL_FILE_PATH}"
+#                 # Return final log, and the file path for Gradio's download component
+#                 return log_output, MODEL_FILE_PATH
+#             else:
+#                 log_output += f"\n⚠️ WARNING: Training completed, but model file not found at expected path ({MODEL_FILE_PATH})."
+#                 return log_output, None
+#         else:
+#             log_output += f"\n\n❌ TRAINING FAILED with return code {return_code}. Check logs above."
+#             return log_output, None
+#
+#     except FileNotFoundError:
+#         return f"❌ ERROR: The training script '{TRAINING_SCRIPT}' was not found. Ensure it is in the root directory of your Space.", None
+#     except Exception as e:
+#         return f"❌ An unexpected error occurred: {e}", None
+#
+#
+# # --- Gradio Interface Setup (using Blocks for a nicer layout) ---
+# with gr.Blocks(title="LayoutLMv3 Fine-Tuning App") as demo:
+#     gr.Markdown("# 🚀 LayoutLMv3 Fine-Tuning on Hugging Face Spaces")
+#     gr.Markdown(
+#         """
+#         Upload your Label Studio JSON file, set your hyperparameters, and click **Train Model** to fine-tune the LayoutLMv3 model using your script.
+#
+#         **Note:** The trained model is saved in the **`checkpoints/`** folder as **`layoutlmv3_crf_passage.pth`**.
+#         """
+#     )
+#
+#     with gr.Row():
+#         with gr.Column(scale=1):
+#             file_input = gr.File(
+#                 label="1. Upload Label Studio JSON Dataset"
+#             )
+#
+#             gr.Markdown("---")
+#             gr.Markdown("### ⚙️ Training Parameters")
+#
+#             batch_size_input = gr.Slider(
+#                 minimum=1, maximum=32, step=1, value=4, label="Batch Size (--batch_size)"
+#             )
+#             epochs_input = gr.Slider(
+#                 minimum=1, maximum=20, step=1, value=5, label="Epochs (--epochs)"
+#             )
+#             lr_input = gr.Number(
+#                 value=5e-5, label="Learning Rate (--lr)"
+#             )
+#             max_len_input = gr.Number(
+#                 value=512, label="Max Sequence Length (--max_len)"
+#             )
+#
+#         with gr.Column(scale=2):
+#             train_button = gr.Button("🔥 Train Model", variant="primary")
+#
+#             log_output = gr.Textbox(
+#                 label="Training Log Output",
+#                 lines=20,
+#                 autoscroll=True,
+#                 placeholder="Click 'Train Model' to start and see real-time logs..."
+#             )
+#
+#             gr.Markdown("---")
+#             gr.Markdown(f"### 🎉 Trained Model Output (Saved to `{MODEL_OUTPUT_DIR}/`)")
+#
+#             # Only providing the download link for the saved .pth model file
+#             model_download = gr.File(label=f"Trained Model File ({MODEL_FILE_NAME})", interactive=False)
+#
+#     # Define the action when the button is clicked
+#     train_button.click(
+#         fn=train_model,
+#         inputs=[file_input, batch_size_input, epochs_input, lr_input, max_len_input],
+#         outputs=[log_output, model_download]
+#     )
+#
+# if __name__ == "__main__":
+#     demo.launch()
+
 
 import gradio as gr
 import subprocess
 import os
 import sys
 from datetime import datetime
+import shutil
 
 # FIX: Update the script name to the correct one you uploaded
 TRAINING_SCRIPT = "HF_LayoutLM_with_Passage.py"
@@ -339,15 +500,13 @@ def train_model(dataset_file: gr.File, batch_size: int, epochs: int, lr: float,
 
     # 2. File Handling: Use the temporary path of the uploaded file
     if dataset_file is None:
-        yield "❌ ERROR: Please upload a file.", None
-        return
+        return "❌ ERROR: Please upload a file.", None
 
     # Using .name (Corrected in previous steps)
     input_path = dataset_file.name
 
     if not input_path.lower().endswith(".json"):
-        yield "❌ ERROR: Please upload a valid Label Studio JSON file (.json).", None
-        return
+        return "❌ ERROR: Please upload a valid Label Studio JSON file (.json).", None
 
     progress(0.1, desc="Starting LayoutLMv3 Training...")
 
@@ -356,7 +515,6 @@ def train_model(dataset_file: gr.File, batch_size: int, epochs: int, lr: float,
     # 3. Construct the subprocess command
     command = [
         sys.executable,
-        # Now uses the corrected TRAINING_SCRIPT variable
         TRAINING_SCRIPT,
         "--mode", "train",
         "--input", input_path,
@@ -367,7 +525,6 @@ def train_model(dataset_file: gr.File, batch_size: int, epochs: int, lr: float,
     ]
 
     log_output += f"Executing command: {' '.join(command)}\n\n"
-    yield log_output, None  # Yield the command to the log output
 
     try:
         # 4. Run the training script and capture output
@@ -382,7 +539,8 @@ def train_model(dataset_file: gr.File, batch_size: int, epochs: int, lr: float,
         # Stream logs in real-time
         for line in iter(process.stdout.readline, ""):
             log_output += line
-            yield log_output, None  # Send partial log to Gradio output
+            # Print to console as well for debugging
+            print(line, end='')
 
         process.stdout.close()
         return_code = process.wait()
@@ -390,81 +548,165 @@ def train_model(dataset_file: gr.File, batch_size: int, epochs: int, lr: float,
         # 5. Check for successful completion
         if return_code == 0:
             log_output += "\n✅ TRAINING COMPLETE! Model saved."
-
-            # 6. Prepare download links based on script's saved path
-            model_exists = os.path.exists(MODEL_FILE_PATH)
-
-            if model_exists:
-                log_output += f"\nModel path: {MODEL_FILE_PATH}"
-                # Return final log, and the file path for Gradio's download component
-                return log_output, MODEL_FILE_PATH
+            print("\n✅ TRAINING COMPLETE! Model saved.")
+
+            # 6. Verify model file exists
+            if os.path.exists(MODEL_FILE_PATH):
+                file_size = os.path.getsize(MODEL_FILE_PATH) / (1024 * 1024)  # Size in MB
+                log_output += f"\n📦 Model file: {MODEL_FILE_PATH}"
+                log_output += f"\n📊 Model size: {file_size:.2f} MB"
+                log_output += f"\n⬇️ Click the download button below to save your model!"
+
+                print(f"\n✅ Model exists at: {MODEL_FILE_PATH} ({file_size:.2f} MB)")
+
+                # Create a copy in the root directory for easier access
+                root_copy = MODEL_FILE_NAME
+                try:
+                    shutil.copy2(MODEL_FILE_PATH, root_copy)
+                    log_output += f"\n📋 Copy created: {root_copy}"
+                    print(f"✅ Created copy at: {root_copy}")
+                except Exception as e:
+                    log_output += f"\n⚠️ Could not create root copy: {e}"
+                    root_copy = MODEL_FILE_PATH
+
+                # Return the full absolute path to ensure Gradio can find it
+                absolute_path = os.path.abspath(root_copy)
+                log_output += f"\n🔗 Download path: {absolute_path}"
+
+                return log_output, absolute_path
             else:
                 log_output += f"\n⚠️ WARNING: Training completed, but model file not found at expected path ({MODEL_FILE_PATH})."
+                log_output += f"\n🔍 Checking directory contents..."
+
+                # List files in checkpoints directory for debugging
+                if os.path.exists(MODEL_OUTPUT_DIR):
+                    files = os.listdir(MODEL_OUTPUT_DIR)
+                    log_output += f"\n📁 Files in {MODEL_OUTPUT_DIR}: {files}"
+                else:
+                    log_output += f"\n❌ Directory {MODEL_OUTPUT_DIR} does not exist!"
+
                 return log_output, None
         else:
             log_output += f"\n\n❌ TRAINING FAILED with return code {return_code}. Check logs above."
             return log_output, None
 
     except FileNotFoundError:
-        return f"❌ ERROR: The training script '{TRAINING_SCRIPT}' was not found. Ensure it is in the root directory of your Space.", None
+        error_msg = f"❌ ERROR: The training script '{TRAINING_SCRIPT}' was not found. Ensure it is in the root directory of your Space."
+        print(error_msg)
+        return error_msg, None
     except Exception as e:
-        return f"❌ An unexpected error occurred: {e}", None
+        error_msg = f"❌ An unexpected error occurred: {e}"
+        print(error_msg)
+        import traceback
+        print(traceback.format_exc())
+        return error_msg, None
 
 
 # --- Gradio Interface Setup (using Blocks for a nicer layout) ---
-with gr.Blocks(title="LayoutLMv3 Fine-Tuning App") as demo:
+with gr.Blocks(title="LayoutLMv3 Fine-Tuning App", theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🚀 LayoutLMv3 Fine-Tuning on Hugging Face Spaces")
     gr.Markdown(
         """
-        Upload your Label Studio JSON file, set your hyperparameters, and click **Train Model** to fine-tune the LayoutLMv3 model using your script.
+        Upload your Label Studio JSON file, set your hyperparameters, and click **Train Model** to fine-tune the LayoutLMv3 model.
+
+        **⚠️ IMPORTANT - Free Tier Users:**
+        - **Download your model IMMEDIATELY** after training completes! 
+        - The model file is **temporary** and will be deleted when the Space restarts.
+        - The download button will appear below once training is complete.
+        - Model is saved as: **`layoutlmv3_crf_passage.pth`**
 
-        **Note:** The trained model is saved in the **`checkpoints/`** folder as **`layoutlmv3_crf_passage.pth`**.
+        **⏱️ Timeout Note:** Training may timeout on free tier. Consider reducing epochs or batch size for faster training.
         """
     )
 
     with gr.Row():
         with gr.Column(scale=1):
+            gr.Markdown("### 📁 Dataset Upload")
             file_input = gr.File(
-                label="1. Upload Label Studio JSON Dataset"
+                label="Upload Label Studio JSON Dataset",
+                file_types=[".json"]
             )
 
             gr.Markdown("---")
             gr.Markdown("### ⚙️ Training Parameters")
 
             batch_size_input = gr.Slider(
-                minimum=1, maximum=32, step=1, value=4, label="Batch Size (--batch_size)"
+                minimum=1, maximum=16, step=1, value=4,
+                label="Batch Size",
+                info="Smaller = less memory, slower training"
             )
             epochs_input = gr.Slider(
-                minimum=1, maximum=20, step=1, value=5, label="Epochs (--epochs)"
+                minimum=1, maximum=10, step=1, value=3,
+                label="Epochs",
+                info="Fewer epochs = faster training (recommended: 3-5)"
             )
             lr_input = gr.Number(
-                value=5e-5, label="Learning Rate (--lr)"
+                value=5e-5, label="Learning Rate",
+                info="Default: 5e-5"
             )
-            max_len_input = gr.Number(
-                value=512, label="Max Sequence Length (--max_len)"
+            max_len_input = gr.Slider(
+                minimum=128, maximum=512, step=128, value=512,
+                label="Max Sequence Length",
+                info="Shorter = faster training, less memory"
             )
 
+            train_button = gr.Button("🔥 Start Training", variant="primary", size="lg")
+
         with gr.Column(scale=2):
-            train_button = gr.Button("🔥 Train Model", variant="primary")
+            gr.Markdown("### 📊 Training Progress")
 
             log_output = gr.Textbox(
-                label="Training Log Output",
-                lines=20,
+                label="Training Logs",
+                lines=25,
+                max_lines=30,
                 autoscroll=True,
-                placeholder="Click 'Train Model' to start and see real-time logs..."
+                show_copy_button=True,
+                placeholder="Click 'Start Training' to begin...\n\nLogs will appear here in real-time."
             )
 
-            gr.Markdown("---")
-            gr.Markdown(f"### 🎉 Trained Model Output (Saved to `{MODEL_OUTPUT_DIR}/`)")
+            gr.Markdown("### ⬇️ Download Trained Model")
 
-            # Only providing the download link for the saved .pth model file
-            model_download = gr.File(label=f"Trained Model File ({MODEL_FILE_NAME})", interactive=False)
+            model_download = gr.File(
+                label="Trained Model File (layoutlmv3_crf_passage.pth)",
+                interactive=False,
+                visible=True
+            )
+
+            gr.Markdown(
+                """
+                **📥 Download Instructions:**
+                1. Wait for training to complete (✅ appears in logs)
+                2. Click the download button/icon that appears above
+                3. Save the `.pth` file to your local machine
+                4. **Do this immediately** - file is temporary!
+
+                **🔧 Troubleshooting:**
+                - If download button doesn't appear, check the logs for errors
+                - Try reducing epochs or batch size if timeout occurs
+                - Ensure your JSON file is properly formatted
+                """
+            )
 
     # Define the action when the button is clicked
     train_button.click(
         fn=train_model,
         inputs=[file_input, batch_size_input, epochs_input, lr_input, max_len_input],
-        outputs=[log_output, model_download]
+        outputs=[log_output, model_download],
+        api_name="train"
+    )
+
+    # Add example info
+    gr.Markdown(
+        """
+        ---
+        ### 📖 About
+        This Space fine-tunes LayoutLMv3 with CRF for document understanding tasks including:
+        - Questions, Options, Answers
+        - Section Headings
+        - Passages
+
+        **Model Details:** LayoutLMv3-base + CRF layer for sequence labeling
+        """
     )
 
 if __name__ == "__main__":