delete cache button

app.py

@@ -66,6 +66,49 @@ def retrieve_model():
         return log_output, None, gr.Button(visible=False)
 
 
+
+
+def clear_memory(dataset_file: gr.File):
+    """
+    Deletes the model output directory and the uploaded dataset file.
+    """
+    MODEL_OUTPUT_DIR = "checkpoints"
+
+    log_output = f"--- Memory Clear Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ---\n"
+
+    # 1. Clear Model Checkpoints Directory
+    if os.path.exists(MODEL_OUTPUT_DIR):
+        try:
+            shutil.rmtree(MODEL_OUTPUT_DIR)
+            log_output += f"✅ Successfully deleted model directory: {MODEL_OUTPUT_DIR}\n"
+        except Exception as e:
+            log_output += f"❌ ERROR deleting model directory {MODEL_OUTPUT_DIR}: {e}\n"
+    else:
+        log_output += f"ℹ️ Model directory not found: {MODEL_OUTPUT_DIR} (Nothing to delete)\n"
+
+    # 2. Clear Uploaded Dataset File (Temporary file cleanup)
+    if dataset_file is not None:
+        input_path = dataset_file.name if hasattr(dataset_file, 'name') else str(dataset_file)
+        if os.path.exists(input_path):
+            try:
+                os.remove(input_path)
+                log_output += f"✅ Successfully deleted uploaded dataset file: {input_path}\n"
+            except Exception as e:
+                log_output += f"❌ ERROR deleting dataset file {input_path}: {e}\n"
+        else:
+            log_output += f"ℹ️ Uploaded dataset file not found at {input_path}.\n"
+    else:
+        log_output += f"ℹ️ No dataset file currently tracked for deletion.\n"
+
+    # 3. Final message and state reset
+    log_output += f"--- Memory Clear Complete: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ---\n"
+    log_output += "✨ Files and checkpoints have been removed. You can now start a fresh training run."
+
+    # Reset log_output, model_path_state, download_btn visibility, and model_download component
+    return log_output, None, gr.Button(visible=False), None
+
+
+
 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.
@@ -280,6 +323,7 @@ with gr.Blocks(title="LayoutLMv3 Fine-Tuning App", theme=gr.themes.Soft()) as de
 
             train_button = gr.Button("🔥 Start Training", variant="primary", size="lg")
             check_button = gr.Button("🔍 Check Model Status/Download", variant="secondary", size="lg")
+            clear_button = gr.Button("🧹 Clear Model/Dataset Files", variant="stop", size="lg")
 
         with gr.Column(scale=2):
             gr.Markdown("### 📊 Training Progress (Real-Time Logs)")
@@ -312,6 +356,8 @@ with gr.Blocks(title="LayoutLMv3 Fine-Tuning App", theme=gr.themes.Soft()) as de
                 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",
@@ -319,6 +365,12 @@ with gr.Blocks(title="LayoutLMv3 Fine-Tuning App", theme=gr.themes.Soft()) as de
                 visible=True
             )
 
+            clear_button.click(
+                fn=clear_memory,
+                inputs=[file_input],  # Pass the uploaded file object to delete the temp file
+                outputs=[log_output, model_path_state, download_btn, model_download]
+            )
+
             gr.Markdown(
                 """
                 **📥 Download Instructions:**