Rename train.py to app.py

app.py

@@ -0,0 +1,369 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+import pdfplumber
+import json
+import os
+import re
+from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Model
+from TorchCRF import CRF
+
+# ---------------------------------------------------------
+# 1. CONFIGURATION
+# ---------------------------------------------------------
+# Ensure this filename matches exactly what you uploaded to the Space
+MODEL_FILENAME = "layoutlmv3_bilstm_crf_hybrid.pth" 
+BASE_MODEL_ID = "microsoft/layoutlmv3-base"
+
+# Define your labels exactly as they were during training
+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()}
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+tokenizer = LayoutLMv3TokenizerFast.from_pretrained(BASE_MODEL_ID)
+
+# ---------------------------------------------------------
+# 2. MODEL ARCHITECTURE
+# ---------------------------------------------------------
+# ⚠️ ACTION REQUIRED: 
+# Replace this class with the exact class definition of your 
+# NEW HYBRID MODEL. The class name and structure must match 
+# what was used when you saved 'layoutlmv3_nonlinear_scratch.pth'.
+# ---------------------------------------------------------
+# ---------------------------------------------------------
+# 2. MODEL ARCHITECTURE (LayoutLMv3 + BiLSTM + CRF)
+# ---------------------------------------------------------
+class HybridModel(nn.Module):
+    def __init__(self, num_labels):
+        super().__init__()
+        self.layoutlm = LayoutLMv3Model.from_pretrained(BASE_MODEL_ID)
+        
+        # Config for BiLSTM
+        hidden_size = self.layoutlm.config.hidden_size # Usually 768
+        lstm_hidden_size = hidden_size // 2  # 384, so bidirectional output is 768
+        
+        # BiLSTM Layer
+        # input_size=768, hidden=384, bidir=True -> output_dim = 384 * 2 = 768
+        self.lstm = nn.LSTM(
+            input_size=hidden_size,
+            hidden_size=lstm_hidden_size,
+            num_layers=1,
+            batch_first=True,
+            bidirectional=True
+        )
+        
+        # Dropout (Optional, check if you used this in training)
+        self.dropout = nn.Dropout(0.1)
+        
+        # Classifier: Maps BiLSTM output (768) to Label count
+        self.classifier = nn.Linear(lstm_hidden_size * 2, num_labels)
+        
+        # CRF Layer
+        self.crf = CRF(num_labels)
+
+    def forward(self, input_ids, bbox, attention_mask, labels=None):
+        # 1. LayoutLMv3 Base
+        outputs = self.layoutlm(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask)
+        sequence_output = outputs.last_hidden_state  # [Batch, Seq, 768]
+        
+        # 2. BiLSTM
+        # LSTM returns (output, (h_n, c_n)). We only need output.
+        lstm_output, _ = self.lstm(sequence_output)  # [Batch, Seq, 768]
+        
+        # 3. Dropout & Classifier
+        lstm_output = self.dropout(lstm_output)
+        emissions = self.classifier(lstm_output)     # [Batch, Seq, Num_Labels]
+
+        # 4. CRF
+        if labels is not None:
+            # Training/Eval (Loss)
+            log_likelihood = self.crf(emissions, labels, mask=attention_mask.bool())
+            return -log_likelihood.mean()
+        else:
+            # Inference (Prediction Tags)
+            return self.crf.viterbi_decode(emissions, mask=attention_mask.bool())
+# ---------------------------------------------------------
+# 3. MODEL LOADING LOGIC
+# ---------------------------------------------------------
+model = None
+
+def load_model():
+    global model
+    if model is None:
+        print(f"🔄 Loading model from {MODEL_FILENAME}...")
+        if not os.path.exists(MODEL_FILENAME):
+            raise FileNotFoundError(f"❌ Model file '{MODEL_FILENAME}' not found. Please upload it to the Files tab of your Space.")
+            
+        # Initialize the model structure
+        model = HybridModel(num_labels=len(LABELS))
+        
+        # Load weights
+        try:
+            state_dict = torch.load(MODEL_FILENAME, map_location=device)
+            model.load_state_dict(state_dict)
+        except RuntimeError as e:
+            raise RuntimeError(f"❌ State dictionary mismatch. Ensure the 'HybridModel' class structure in app.py matches the model you trained.\nDetails: {e}")
+
+        model.to(device)
+        model.eval()
+        print("✅ Model loaded successfully.")
+    return model
+
+# ---------------------------------------------------------
+# 4. JSON CONVERSION LOGIC (Your Custom Logic)
+# ---------------------------------------------------------
+def convert_bio_to_structured_json(predictions):
+    structured_data = []
+    current_item = None
+    current_option_key = None
+    current_passage_buffer = []
+    current_text_buffer = []
+    first_question_started = False
+    last_entity_type = None
+    just_finished_i_option = False
+    is_in_new_passage = False
+
+    def finalize_passage_to_item(item, passage_buffer):
+        if passage_buffer:
+            passage_text = re.sub(r'\s{2,}', ' ', ' '.join(passage_buffer)).strip()
+            if item.get('passage'): item['passage'] += ' ' + passage_text
+            else: item['passage'] = passage_text
+        passage_buffer.clear()
+
+    # Flatten predictions list if strictly page-separated
+    flat_predictions = []
+    for page in predictions:
+        flat_predictions.extend(page['data'])
+
+    for idx, item in enumerate(flat_predictions):
+        word = item['word']
+        label = item['predicted_label']
+        entity_type = label[2:].strip() if label.startswith(('B-', 'I-')) else None
+        current_text_buffer.append(word)
+        
+        previous_entity_type = last_entity_type
+        is_passage_label = (entity_type == 'PASSAGE')
+
+        if not first_question_started:
+            if label != 'B-QUESTION' and not is_passage_label:
+                just_finished_i_option = False
+                is_in_new_passage = False
+                continue
+            if is_passage_label:
+                current_passage_buffer.append(word)
+                last_entity_type = 'PASSAGE'
+                just_finished_i_option = False
+                is_in_new_passage = False
+                continue
+
+        if label == 'B-QUESTION':
+            if not first_question_started:
+                header_text = ' '.join(current_text_buffer[:-1]).strip()
+                if header_text or current_passage_buffer:
+                    metadata_item = {'type': 'METADATA', 'passage': ''}
+                    finalize_passage_to_item(metadata_item, current_passage_buffer)
+                    if header_text: metadata_item['text'] = header_text
+                    structured_data.append(metadata_item)
+                first_question_started = True
+                current_text_buffer = [word]
+
+            if current_item is not None:
+                finalize_passage_to_item(current_item, current_passage_buffer)
+                current_item['text'] = ' '.join(current_text_buffer[:-1]).strip()
+                structured_data.append(current_item)
+                current_text_buffer = [word]
+
+            current_item = {
+                'question': word, 'options': {}, 'answer': '', 'passage': '', 'text': ''
+            }
+            current_option_key = None
+            last_entity_type = 'QUESTION'
+            just_finished_i_option = False
+            is_in_new_passage = False
+            continue
+
+        if current_item is not None:
+            if is_in_new_passage:
+                if 'new_passage' not in current_item: current_item['new_passage'] = word
+                else: current_item['new_passage'] += f' {word}'
+                if label.startswith('B-') or (label.startswith('I-') and entity_type != 'PASSAGE'):
+                    is_in_new_passage = False
+                if label.startswith(('B-', 'I-')): last_entity_type = entity_type
+                continue
+
+            is_in_new_passage = False
+            if label.startswith('B-'):
+                if entity_type in ['QUESTION', 'OPTION', 'ANSWER', 'SECTION_HEADING']:
+                    finalize_passage_to_item(current_item, current_passage_buffer)
+                    current_passage_buffer = []
+                last_entity_type = entity_type
+                if entity_type == 'PASSAGE':
+                    if previous_entity_type == 'OPTION' and just_finished_i_option:
+                        current_item['new_passage'] = word
+                        is_in_new_passage = True
+                    else: current_passage_buffer.append(word)
+                elif entity_type == 'OPTION':
+                    current_option_key = word
+                    current_item['options'][current_option_key] = word
+                    just_finished_i_option = False
+                elif entity_type == 'ANSWER':
+                    current_item['answer'] = word
+                    current_option_key = None
+                    just_finished_i_option = False
+                elif entity_type == 'QUESTION':
+                    current_item['question'] += f' {word}'
+                    just_finished_i_option = False
+
+            elif label.startswith('I-'):
+                if entity_type == 'QUESTION': current_item['question'] += f' {word}'
+                elif entity_type == 'PASSAGE':
+                    if previous_entity_type == 'OPTION' and just_finished_i_option:
+                        current_item['new_passage'] = word
+                        is_in_new_passage = True
+                    else:
+                        if not current_passage_buffer: last_entity_type = 'PASSAGE'
+                        current_passage_buffer.append(word)
+                elif entity_type == 'OPTION' and current_option_key is not None:
+                    current_item['options'][current_option_key] += f' {word}'
+                    just_finished_i_option = True
+                elif entity_type == 'ANSWER': current_item['answer'] += f' {word}'
+                just_finished_i_option = (entity_type == 'OPTION')
+
+    if current_item is not None:
+        finalize_passage_to_item(current_item, current_passage_buffer)
+        current_item['text'] = ' '.join(current_text_buffer).strip()
+        structured_data.append(current_item)
+
+    # Final Cleanup
+    for item in structured_data:
+        if 'text' in item: item['text'] = re.sub(r'\s{2,}', ' ', item['text']).strip()
+        if 'new_passage' in item: item['new_passage'] = re.sub(r'\s{2,}', ' ', item['new_passage']).strip()
+
+    return structured_data
+
+# ---------------------------------------------------------
+# 5. INFERENCE PIPELINE
+# ---------------------------------------------------------
+def process_pdf(pdf_file):
+    if pdf_file is None:
+        return None, "⚠️ Please upload a PDF file."
+
+    try:
+        active_model = load_model()
+        
+        # A. Extract Text and Boxes
+        extracted_pages = []
+        with pdfplumber.open(pdf_file.name) as pdf:
+            for page_idx, page in enumerate(pdf.pages):
+                width, height = page.width, page.height
+                words_data = page.extract_words()
+                
+                page_tokens = []
+                page_bboxes = []
+                
+                for w in words_data:
+                    text = w['text']
+                    # Normalize bbox to 0-1000 scale
+                    x0 = int((w['x0'] / width) * 1000)
+                    top = int((w['top'] / height) * 1000)
+                    x1 = int((w['x1'] / width) * 1000)
+                    bottom = int((w['bottom'] / height) * 1000)
+                    
+                    # Safety clamp
+                    box = [max(0, min(x0, 1000)), max(0, min(top, 1000)), 
+                           max(0, min(x1, 1000)), max(0, min(bottom, 1000))]
+                    
+                    page_tokens.append(text)
+                    page_bboxes.append(box)
+                extracted_pages.append({"page_id": page_idx, "tokens": page_tokens, "bboxes": page_bboxes})
+
+        # B. Run Inference
+        raw_predictions = []
+        for page in extracted_pages:
+            tokens = page['tokens']
+            bboxes = page['bboxes']
+            if not tokens: continue
+
+            # Tokenize
+            encoding = tokenizer(
+                tokens, 
+                boxes=bboxes, 
+                return_tensors="pt", 
+                padding="max_length", 
+                truncation=True, 
+                max_length=512,
+                return_offsets_mapping=True
+            )
+            
+            input_ids = encoding.input_ids.to(device)
+            bbox = encoding.bbox.to(device)
+            attention_mask = encoding.attention_mask.to(device)
+            
+            # Predict
+            with torch.no_grad():
+                # NOTE: If your hybrid model requires 'pixel_values', 
+                # you will need to add image extraction logic above and pass it here.
+                preds = active_model(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask)
+                
+                # Check if preds returns a tuple (loss, tags) or just tags
+                # The CRF implementation usually returns a list of lists of tags in viterbi_decode
+                pred_tags = preds[0] if isinstance(preds, tuple) else preds[0] 
+                # Note: Standard CRF.viterbi_decode returns List[List[int]], so [0] gets the first batch item
+
+            # Alignment
+            word_ids = encoding.word_ids()
+            aligned_data = []
+            prev_word_idx = None
+            
+            for i, word_idx in enumerate(word_ids):
+                if word_idx is None: continue
+                if word_idx != prev_word_idx:
+                    # pred_tags is likely a list of ints. 
+                    # If pred_tags[i] fails, your max_length might be cutting off tags, 
+                    # or the model output shape differs from the token length.
+                    if i < len(pred_tags): 
+                        label_id = pred_tags[i]
+                        label_str = ID2LABEL.get(label_id, "O")
+                        aligned_data.append({"word": tokens[word_idx], "predicted_label": label_str})
+                prev_word_idx = word_idx
+            raw_predictions.append({"data": aligned_data})
+
+        # C. Convert to Structured JSON
+        final_json = convert_bio_to_structured_json(raw_predictions)
+        
+        # Save output
+        output_filename = "structured_output.json"
+        with open(output_filename, "w", encoding="utf-8") as f:
+            json.dump(final_json, f, indent=2, ensure_ascii=False)
+            
+        return output_filename, f"✅ Success! Processed {len(extracted_pages)} pages. Extracted {len(final_json)} items."
+
+    except Exception as e:
+        import traceback
+        return None, f"❌ Error:\n{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+
+# ---------------------------------------------------------
+# 6. GRADIO INTERFACE
+# ---------------------------------------------------------
+iface = gr.Interface(
+    fn=process_pdf,
+    inputs=gr.File(label="Upload PDF", file_types=[".pdf"]),
+    outputs=[
+        gr.File(label="Download JSON Output"),
+        gr.Textbox(label="Status Log", lines=10)
+    ],
+    title="Hybrid Model Inference: PDF to JSON",
+    description="Upload a document to extract structured data using the custom Hybrid LayoutLMv3 model.",
+    flagging_mode="never"
+)
+
+if __name__ == "__main__":
+    iface.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)