Rename train.py to app.py

app.py

@@ -0,0 +1,264 @@
+import gradio as gr
+import subprocess
+import os
+import sys
+from datetime import datetime
+import shutil
+
+# --- CONFIGURATION UPDATED FOR HYBRID MODEL ---
+TRAINING_SCRIPT = "train_hybrid.py"
+MODEL_OUTPUT_DIR = "checkpoints"
+MODEL_FILE_NAME = "layoutlmv3_bilstm_crf_hybrid.pth"
+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.
+    """
+    if os.path.exists(MODEL_FILE_PATH):
+        file_size = os.path.getsize(MODEL_FILE_PATH) / (1024 * 1024)  # Size in MB
+
+        # 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"hybrid_model_recovered_{timestamp}.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! The Hybrid LayoutLMv3+BiLSTM+CRF model was found.\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 clear_memory(dataset_file: gr.File):
+    """
+    Deletes the model output directory and the uploaded dataset file.
+    """
+    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"
+
+    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."
+
+    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.
+    """
+    # 1. Setup: Create output directory
+    os.makedirs(MODEL_OUTPUT_DIR, exist_ok=True)
+
+    # 2. File Handling
+    if dataset_file is None:
+        yield "❌ ERROR: Please upload a file.", None, gr.Button(visible=False)
+        return
+
+    input_path = dataset_file.name if hasattr(dataset_file, 'name') else str(dataset_file)
+
+    if not os.path.exists(input_path):
+        yield f"❌ ERROR: Uploaded file not found at {input_path}.", None, gr.Button(visible=False)
+        return
+
+    progress(0.1, desc="Initializing Hybrid Model Training...")
+
+    log_output = f"--- Training Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ---\n"
+    log_output += f"🤖 Architecture: LayoutLMv3 + BiLSTM + CRF\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, gr.Button(visible=False)
+
+    try:
+        # 4. Run the training script
+        process = subprocess.Popen(
+            command,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            text=True,
+            bufsize=1
+        )
+
+        # Stream logs
+        for line in iter(process.stdout.readline, ""):
+            log_output += line
+            print(line, end='')
+            yield log_output, None, gr.Button(visible=False)
+
+        process.stdout.close()
+        return_code = process.wait()
+
+        # 5. Check completion
+        if return_code == 0:
+            log_output += "\n" + "=" * 60 + "\n"
+            log_output += "✅ HYBRID TRAINING COMPLETE!\n"
+            log_output += "=" * 60 + "\n"
+
+            if os.path.exists(MODEL_FILE_PATH):
+                file_size = os.path.getsize(MODEL_FILE_PATH) / (1024 * 1024)
+                log_output += f"\n📦 Model file found: {MODEL_FILE_PATH} ({file_size:.2f} MB)"
+
+                # Copy for download
+                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"hybrid_model_{timestamp}.pth")
+
+                try:
+                    shutil.copy2(MODEL_FILE_PATH, temp_model_path)
+                    download_path = temp_model_path
+                except Exception as e:
+                    log_output += f"\n⚠️ Copy failed: {e}, using original path"
+                    download_path = MODEL_FILE_PATH
+
+                log_output += f"\n\n⬇️ Click the '📥 Download Model' button below."
+                yield log_output, download_path, gr.Button(visible=True)
+                return
+            else:
+                log_output += f"\n❌ Error: Training finished but {MODEL_FILE_PATH} was not found."
+                yield log_output, None, gr.Button(visible=False)
+                return
+        else:
+            log_output += f"\n❌ TRAINING FAILED with return code {return_code}\n"
+            yield log_output, None, gr.Button(visible=False)
+            return
+
+    except FileNotFoundError:
+        yield log_output + f"\n❌ ERROR: '{TRAINING_SCRIPT}' not found.", None, gr.Button(visible=False)
+    except Exception as e:
+        yield log_output + f"\n❌ Unexpected Error: {e}", None, gr.Button(visible=False)
+
+
+# --- Gradio Interface Setup ---
+with gr.Blocks(title="Hybrid LayoutLM Training", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🧬 Hybrid LayoutLMv3 + BiLSTM + CRF Training")
+    gr.Markdown(
+        """
+        **Architecture:** This app trains a state-of-the-art stack:
+        1. **LayoutLMv3** (Visual & Textual Embeddings)
+        2. **Bi-LSTM** (Sequence Context Modeling)
+        3. **CRF** (Label Consistency Enforcement)
+        
+        **Instructions:** Upload your Label Studio JSON, set parameters, and train. 
+        **Note:** This model is slower to train than standard LayoutLM but typically achieves higher accuracy on complex layouts.
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### 📁 Dataset")
+            file_input = gr.File(label="Upload Label Studio JSON", file_types=[".json"])
+
+            gr.Markdown("### ⚙️ Hyperparameters")
+            batch_size_input = gr.Slider(1, 16, value=4, step=1, label="Batch Size")
+            epochs_input = gr.Slider(1, 10, value=5, step=1, label="Epochs")
+            lr_input = gr.Number(value=2e-5, label="Learning Rate (Backbone)", info="LSTM/CRF head uses 1e-4")
+            max_len_input = gr.Slider(128, 512, value=512, step=128, label="Max Seq Len")
+
+            train_button = gr.Button("🔥 Start Hybrid Training", variant="primary", size="lg")
+            check_button = gr.Button("🔍 Check Status / Recover Model", variant="secondary")
+            clear_button = gr.Button("🧹 Clear Files", variant="stop")
+
+        with gr.Column(scale=2):
+            log_output = gr.Textbox(
+                label="Training Logs", lines=25, autoscroll=True, show_copy_button=True,
+                placeholder="Logs will appear here..."
+            )
+            
+            download_btn = gr.Button("📥 Download Hybrid Model", variant="primary", size="lg", visible=False)
+            
+            # State and hidden download component
+            model_path_state = gr.State(value=None)
+            model_download = gr.File(label="Download", interactive=False, visible=True)
+
+    # Actions
+    train_button.click(
+        fn=train_model,
+        inputs=[file_input, batch_size_input, epochs_input, lr_input, max_len_input],
+        outputs=[log_output, model_path_state, download_btn]
+    )
+
+    check_button.click(
+        fn=retrieve_model,
+        inputs=[],
+        outputs=[log_output, model_path_state, download_btn]
+    )
+
+    clear_button.click(
+        fn=clear_memory,
+        inputs=[file_input],
+        outputs=[log_output, model_path_state, download_btn, model_download]
+    )
+
+    download_btn.click(
+        fn=lambda path: path,
+        inputs=[model_path_state],
+        outputs=[model_download]
+    )
+
+if __name__ == "__main__":
+    demo.launch()

train.py

@@ -1,244 +0,0 @@
-import json
-import argparse
-import os
-import random
-import torch
-import torch.nn as nn
-from torch.utils.data import Dataset, DataLoader, random_split
-from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Model
-from TorchCRF import CRF
-from torch.optim import AdamW
-from tqdm import tqdm
-from sklearn.metrics import precision_recall_fscore_support
-
-# --- Configuration ---
-MAX_BBOX_DIMENSION = 1000
-MAX_SHIFT = 30
-AUGMENTATION_FACTOR = 1
-BASE_MODEL_ID = "heerjtdev/MLP_LayoutLM"
-
-# -------------------------
-# Step 1: Preprocessing
-# -------------------------
-def preprocess_labelstudio(input_path, output_path):
-    with open(input_path, "r", encoding="utf-8") as f:
-        data = json.load(f)
-
-    processed = []
-    print(f"🔄 Starting preprocessing of {len(data)} documents...")
-
-    for item in data:
-        words = item["data"]["original_words"]
-        bboxes = item["data"]["original_bboxes"]
-        labels = ["O"] * len(words)
-
-        clamped_bboxes = []
-        for bbox in bboxes:
-            x_min, y_min, x_max, y_max = bbox
-            new_x_min = max(0, min(x_min, 1000))
-            new_y_min = max(0, min(y_min, 1000))
-            new_x_max = max(0, min(x_max, 1000))
-            new_y_max = max(0, min(y_max, 1000))
-            if new_x_min > new_x_max: new_x_min = new_x_max
-            if new_y_min > new_y_max: new_y_min = new_y_max
-            clamped_bboxes.append([new_x_min, new_y_min, new_x_max, new_y_max])
-
-        if "annotations" in item:
-            for ann in item["annotations"]:
-                for res in ann["result"]:
-                    if "value" in res and "labels" in res["value"]:
-                        text = res["value"]["text"]
-                        tag = res["value"]["labels"][0]
-                        text_tokens = text.split()
-                        for i in range(len(words) - len(text_tokens) + 1):
-                            if words[i:i + len(text_tokens)] == text_tokens:
-                                labels[i] = f"B-{tag}"
-                                for j in range(1, len(text_tokens)):
-                                    labels[i + j] = f"I-{tag}"
-                                break
-
-        processed.append({"tokens": words, "labels": labels, "bboxes": clamped_bboxes})
-
-    with open(output_path, "w", encoding="utf-8") as f:
-        json.dump(processed, f, indent=2, ensure_ascii=False)
-    return output_path
-
-# -------------------------
-# Step 1.5: Augmentation
-# -------------------------
-def translate_bbox(bbox, shift_x, shift_y):
-    x_min, y_min, x_max, y_max = bbox
-    new_x_min = max(0, min(x_min + shift_x, 1000))
-    new_y_min = max(0, min(y_min + shift_y, 1000))
-    new_x_max = max(0, min(x_max + shift_x, 1000))
-    new_y_max = max(0, min(y_max + shift_y, 1000))
-    return [new_x_min, new_y_min, new_x_max, new_y_max]
-
-def augment_sample(sample):
-    shift_x = random.randint(-MAX_SHIFT, MAX_SHIFT)
-    shift_y = random.randint(-MAX_SHIFT, MAX_SHIFT)
-    new_sample = sample.copy()
-    new_sample["bboxes"] = [translate_bbox(b, shift_x, shift_y) for b in sample["bboxes"]]
-    return new_sample
-
-def augment_and_save_dataset(input_json_path, output_json_path):
-    with open(input_json_path, 'r', encoding="utf-8") as f:
-        training_data = json.load(f)
-    augmented_data = []
-    for original_sample in training_data:
-        augmented_data.append(original_sample)
-        for _ in range(AUGMENTATION_FACTOR):
-            augmented_data.append(augment_sample(original_sample))
-    with open(output_json_path, 'w', encoding="utf-8") as f:
-        json.dump(augmented_data, f, indent=2, ensure_ascii=False)
-    return output_json_path
-
-# -------------------------
-# Step 2: Dataset Class
-# -------------------------
-class LayoutDataset(Dataset):
-    def __init__(self, json_path, tokenizer, label2id, max_len=512):
-        with open(json_path, "r", encoding="utf-8") as f:
-            self.data = json.load(f)
-        self.tokenizer = tokenizer
-        self.label2id = label2id
-        self.max_len = max_len
-
-    def __len__(self):
-        return len(self.data)
-
-    def __getitem__(self, idx):
-        item = self.data[idx]
-        words, bboxes, labels = item["tokens"], item["bboxes"], item["labels"]
-        encodings = self.tokenizer(words, boxes=bboxes, padding="max_length", truncation=True, max_length=self.max_len, return_tensors="pt")
-        word_ids = encodings.word_ids(batch_index=0)
-        label_ids = []
-        for word_id in word_ids:
-            if word_id is None:
-                label_ids.append(self.label2id["O"])
-            else:
-                label_ids.append(self.label2id.get(labels[word_id], self.label2id["O"]))
-        encodings["labels"] = torch.tensor(label_ids)
-        return {key: val.squeeze(0) for key, val in encodings.items()}
-
-# -------------------------
-# Step 3: Model Architecture (Non-Linear Head)
-# -------------------------
-
-class LayoutLMv3CRF(nn.Module):
-    def __init__(self, num_labels):
-        super().__init__()
-        # Initializing from scratch (Base weights only)
-        print(f"🔄 Initializing backbone from {BASE_MODEL_ID}...")
-        self.layoutlm = LayoutLMv3Model.from_pretrained(BASE_MODEL_ID)
-
-        hidden_size = self.layoutlm.config.hidden_size
-
-        # NON-LINEAR MLP HEAD
-        # Replacing the simple Linear layer with a deeper architecture
-        self.classifier = nn.Sequential(
-            nn.Linear(hidden_size, hidden_size),
-            nn.GELU(),                  # Non-linear activation
-            nn.LayerNorm(hidden_size),   # Stability for training from scratch
-            nn.Dropout(0.1),
-            nn.Linear(hidden_size, num_labels)
-        )
-
-        self.crf = CRF(num_labels)
-
-    def forward(self, input_ids, bbox, attention_mask, labels=None):
-        outputs = self.layoutlm(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask)
-        sequence_output = outputs.last_hidden_state
-
-        # Pass through the new non-linear head
-        emissions = self.classifier(sequence_output)
-
-        if labels is not None:
-            log_likelihood = self.crf(emissions, labels, mask=attention_mask.bool())
-            return -log_likelihood.mean()
-        else:
-            return self.crf.viterbi_decode(emissions, mask=attention_mask.bool())
-
-# -------------------------
-# Step 4: Training + Evaluation
-# -------------------------
-def train_one_epoch(model, dataloader, optimizer, device):
-    model.train()
-    total_loss = 0
-    for batch in tqdm(dataloader, desc="Training"):
-        batch = {k: v.to(device) for k, v in batch.items()}
-        labels = batch.pop("labels")
-        optimizer.zero_grad()
-        loss = model(**batch, labels=labels)
-        loss.backward()
-        optimizer.step()
-        total_loss += loss.item()
-    return total_loss / len(dataloader)
-
-def evaluate(model, dataloader, device, id2label):
-    model.eval()
-    all_preds, all_labels = [], []
-    with torch.no_grad():
-        for batch in tqdm(dataloader, desc="Evaluating"):
-            batch = {k: v.to(device) for k, v in batch.items()}
-            labels = batch.pop("labels").cpu().numpy()
-            preds = model(**batch)
-            for p, l, mask in zip(preds, labels, batch["attention_mask"].cpu().numpy()):
-                valid = mask == 1
-                l_valid = l[valid].tolist()
-                all_labels.extend(l_valid)
-                all_preds.extend(p[:len(l_valid)])
-    precision, recall, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average="micro", zero_division=0)
-    return precision, recall, f1
-
-# -------------------------
-# Step 5: Main Execution
-# -------------------------
-def main(args):
-    labels = ["O", "B-QUESTION", "I-QUESTION", "B-OPTION", "I-OPTION", "B-ANSWER", "I-ANSWER", "B-SECTION_HEADING", "I-SECTION_HEADING", "B-PASSAGE", "I-PASSAGE"]
-    label2id = {l: i for i, l in enumerate(labels)}
-    id2label = {i: l for l, i in label2id.items()}
-
-    TEMP_DIR = "temp_intermediate_files"
-    os.makedirs(TEMP_DIR, exist_ok=True)
-
-    # 1. Preprocess & Augment
-    initial_json = os.path.join(TEMP_DIR, "data_bio.json")
-    preprocess_labelstudio(args.input, initial_json)
-    augmented_json = os.path.join(TEMP_DIR, "data_aug.json")
-    final_data_path = augment_and_save_dataset(initial_json, augmented_json)
-
-    # 2. Setup Data
-    tokenizer = LayoutLMv3TokenizerFast.from_pretrained(BASE_MODEL_ID)
-    dataset = LayoutDataset(final_data_path, tokenizer, label2id, max_len=args.max_len)
-    val_size = int(0.2 * len(dataset))
-    train_dataset, val_dataset = random_split(dataset, [len(dataset) - val_size, val_size])
-
-    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
-    val_loader = DataLoader(val_dataset, batch_size=args.batch_size)
-
-    # 3. Model
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = LayoutLMv3CRF(num_labels=len(labels)).to(device)
-    optimizer = AdamW(model.parameters(), lr=args.lr)
-
-    # 4. Loop
-    for epoch in range(args.epochs):
-        loss = train_one_epoch(model, train_loader, optimizer, device)
-        p, r, f1 = evaluate(model, val_loader, device, id2label)
-        print(f"Epoch {epoch+1} | Loss: {loss:.4f} | F1: {f1:.3f}")
-
-        ckpt_path = "checkpoints/layoutlmv3_nonlinear_scratch.pth"
-        os.makedirs("checkpoints", exist_ok=True)
-        torch.save(model.state_dict(), ckpt_path)
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--mode", type=str, default="train")
-    parser.add_argument("--input", type=str, required=True)
-    parser.add_argument("--batch_size", type=int, default=4)
-    parser.add_argument("--epochs", type=int, default=10) # Increased for scratch training
-    parser.add_argument("--lr", type=float, default=2e-5)
-    parser.add_argument("--max_len", type=int, default=512)
-    args = parser.parse_args()
-    main(args)