Spaces:

heerjtdev
/

LayoutLM_train

Running

File size: 26,751 Bytes

# import gradio as gr
# import subprocess
# import os
# import sys
# from datetime import datetime
#
# # The name of your existing training script
# TRAINING_SCRIPT = "LayoutLM_Train_Passage.py"
#
# # --- CORRECTED MODEL PATH BASED ON LayoutLM_Train_Passage.py ---
# 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 or not dataset_file.path.endswith(".json"):
#     #     return "❌ ERROR: Please upload a valid Label Studio JSON file.", None
#
#     input_path = dataset_file.path
#
#     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,
#         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"
#
#     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(server_port=7860, server_name="0.0.0.0")


# import gradio as gr
# import subprocess
# import os
# import sys
# from datetime import datetime
#
# # The name of your existing training script
# TRAINING_SCRIPT = "LayoutLM_Train_Passage.py"
#
# # --- CORRECTED MODEL PATH BASED ON LayoutLM_Train_Passage.py ---
# 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
#
#     # FIX: Gradio returns the path in the .name attribute, not .path
#     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,
#         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__":
#     # Removed server_port and server_name as they are often unnecessary
#     # 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"

# --- 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:
        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"):
        return "❌ ERROR: Please upload a valid Label Studio JSON file (.json).", None

    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,
        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"

    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
            # Print to console as well for debugging
            print(line, end='')

        process.stdout.close()
        return_code = process.wait()

        # 5. Check for successful completion
        if return_code == 0:
            log_output += "\n✅ TRAINING COMPLETE! Model saved."
            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:
        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:
        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", 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.

        **⚠️ 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`**

        **⏱️ 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="Upload Label Studio JSON Dataset",
                file_types=[".json"]
            )

            gr.Markdown("---")
            gr.Markdown("### ⚙️ Training Parameters")

            batch_size_input = gr.Slider(
                minimum=1, maximum=16, step=1, value=4,
                label="Batch Size",
                info="Smaller = less memory, slower training"
            )
            epochs_input = gr.Slider(
                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",
                info="Default: 5e-5"
            )
            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):
            gr.Markdown("### 📊 Training Progress")

            log_output = gr.Textbox(
                label="Training Logs",
                lines=25,
                max_lines=30,
                autoscroll=True,
                show_copy_button=True,
                placeholder="Click 'Start Training' to begin...\n\nLogs will appear here in real-time."
            )

            gr.Markdown("### ⬇️ Download Trained Model")

            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],
        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__":
    demo.launch()