refactor: create predict_pipeline.py and utils.py to modularize app logic

- Created predict_pipeline.py to handle enrichment and inference pipeline
- Added run_prediction_pipeline() with mocked summarization and translation
- Handles uploaded WhatsApp ZIP, merges with description, and runs prediction

- Created utils.py for shared logic used across modules
- Moved AbuseDataset class and label_row_soft function
- Added map_to_3_classes, convert_to_label_strings, and tune_thresholds
- Centralized label_map for consistent mapping

- Updated Gradio UI to import from predict_pipeline
- Improves modularity, reusability, and future maintainability

app.py

@@ -1,20 +1,39 @@
 import gradio as gr
-from train_abuse_model import run_training, evaluate_saved_model, push_model_to_hub
+from train_abuse_model import (
+    run_training,
+    evaluate_saved_model,
+    push_model_to_hub
+)
+from predict_pipeline import run_prediction_pipeline
+
 
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 Abuse Detection Fine-Tuning App")
+    gr.Markdown("## 🧠 Abuse Detection App")
     gr.Markdown("⚠️ Keep this tab open while training or evaluating.")
 
-    with gr.Row():
-        start_btn = gr.Button("🚀 Start Training")
-        eval_btn = gr.Button("🔍 Evaluate Trained Model")
-        push_btn = gr.Button("📤 Push Model to Hub")
+    with gr.Tab("🧪 Train / Evaluate"):
+        with gr.Row():
+            start_btn = gr.Button("🚀 Start Training")
+            eval_btn = gr.Button("🔍 Evaluate Trained Model")
+            push_btn = gr.Button("📤 Push Model to Hub")
+        output_box = gr.Textbox(label="Logs", lines=25, interactive=False)
+        start_btn.click(fn=run_training, outputs=output_box)
+        eval_btn.click(fn=evaluate_saved_model, outputs=output_box)
+        push_btn.click(fn=push_model_to_hub, outputs=output_box)
+
+    with gr.Tab("🔮 Abuse Detection"):
+        desc_input = gr.Textbox(label="📝 Relationship Description", lines=5, placeholder="Write a relationship story here...")
+        chat_upload = gr.File(label="📁 Optional: WhatsApp Chat ZIP (.zip)", file_types=[".zip"])
+        predict_btn = gr.Button("Run Prediction")
 
-    output_box = gr.Textbox(label="Logs", lines=25, interactive=False)
+        enriched_output = gr.Textbox(label="📎 Enriched Input (Used for Prediction)", lines=8, interactive=False)
+        label_output = gr.Textbox(label="🏷️ Predicted Labels", lines=2, interactive=False)
 
-    start_btn.click(fn=run_training, outputs=output_box)
-    eval_btn.click(fn=evaluate_saved_model, outputs=output_box)
-    push_btn.click(fn=push_model_to_hub, outputs=output_box)
+        predict_btn.click(
+            fn=run_prediction_pipeline,
+            inputs=[desc_input, chat_upload],
+            outputs=[enriched_output, label_output]
+        )
 
 if __name__ == "__main__":
     demo.launch()

train_abuse_model.py

@@ -3,6 +3,7 @@
 
 import logging
 import io
+import os
 import time
 import gradio as gr  # ✅ required for progress bar
 from pathlib import Path
@@ -31,6 +32,16 @@ from transformers import (
     TrainingArguments
 )
 
+from utils import (
+    map_to_3_classes,
+    convert_to_label_strings,
+    tune_thresholds,
+    label_map,
+    label_row_soft,
+    AbuseDataset
+)
+
+
 PERSIST_DIR = Path("/home/user/app")
 MODEL_DIR = PERSIST_DIR / "saved_model"
 LOG_FILE = PERSIST_DIR / "training.log"
@@ -56,82 +67,6 @@ logger.info("Transformers version: %s", torch.__version__)
 logger.info("torch.cuda.is_available(): %s", torch.cuda.is_available())
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Label mapping for evaluation
-label_map = {0.0: "no", 0.5: "plausibly", 1.0: "yes"}
-
-# Custom Dataset class
-
-class AbuseDataset(Dataset):
-    def __init__(self, texts, labels, tokenizer):
-        self.encodings = tokenizer(texts, truncation=True, padding=True, max_length=512)
-        self.labels = labels
-
-    def __len__(self):
-        return len(self.labels)
-    
-    def __getitem__(self, idx):
-        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
-        item["labels"] = torch.tensor(self.labels[idx], dtype=torch.float)
-        return item
-    def __getitem__(self, idx):
-        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
-        item["labels"] = torch.tensor(self.labels[idx], dtype=torch.float)
-        return item
-
-
-#  Convert label values to soft scores: "yes" = 1.0, "plausibly" = 0.5, others = 0.0
-def label_row_soft(row):
-    labels = []
-    for col in label_columns:
-        val = str(row[col]).strip().lower()
-        if val == "yes":
-            labels.append(1.0)
-        elif val == "plausibly":
-            labels.append(0.5)
-        else:
-            labels.append(0.0)
-    return labels
-    
-# Function to map probabilities to 3 classes
-# (0.0, 0.5, 1.0) based on thresholds
-def map_to_3_classes(prob_array, low, high):
-    """Map probabilities to 0.0, 0.5, 1.0 using thresholds."""
-    mapped = np.zeros_like(prob_array)
-    mapped[(prob_array > low) & (prob_array <= high)] = 0.5
-    mapped[prob_array > high] = 1.0
-    return mapped
-
-def convert_to_label_strings(array):
-    """Convert float label array to list of strings."""
-    return [label_map[val] for val in array.flatten()]
-
-def tune_thresholds(probs, true_labels, verbose=True):
-    """Search for best (low, high) thresholds by macro F1 score."""
-    best_macro_f1 = 0.0
-    best_low, best_high = 0.0, 0.0
-
-    for low in np.arange(0.2, 0.5, 0.05):
-        for high in np.arange(0.55, 0.8, 0.05):
-            if high <= low:
-                continue
-
-            pred_soft = map_to_3_classes(probs, low, high)
-            pred_str = convert_to_label_strings(pred_soft)
-            true_str = convert_to_label_strings(true_labels)
-
-            _, _, f1, _ = precision_recall_fscore_support(
-                true_str, pred_str,
-                labels=["no", "plausibly", "yes"],
-                average="macro",
-                zero_division=0
-            )
-            if verbose:
-                logger.info(f"low={low:.2f}, high={high:.2f} -> macro F1={f1:.3f}")
-            if f1 > best_macro_f1:
-                best_macro_f1 = f1
-                best_low, best_high = low, high
-
-    return best_low, best_high, best_macro_f1
 
 def evaluate_model_with_thresholds(trainer, test_dataset):
     """Run full evaluation with automatic threshold tuning."""
@@ -198,8 +133,6 @@ def evaluate_saved_model(progress=gr.Progress(track_tqdm=True)):
             eval_dataset=test_dataset
         )
 
-        label_map = {0.0: "no", 0.5: "plausibly", 1.0: "yes"}
-
         # Re-yield from generator
         for line in evaluate_model_with_thresholds(trainer, test_dataset):
             yield line

utils.py

@@ -0,0 +1,85 @@
+import numpy as np
+from sklearn.metrics import precision_recall_fscore_support
+import torch
+from torch.utils.data import Dataset
+
+# Custom Dataset class
+class AbuseDataset(Dataset):
+    def __init__(self, texts, labels, tokenizer):
+        self.encodings = tokenizer(texts, truncation=True, padding=True, max_length=512)
+        self.labels = labels
+
+    def __len__(self):
+        return len(self.labels)
+    
+    def __getitem__(self, idx):
+        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
+        item["labels"] = torch.tensor(self.labels[idx], dtype=torch.float)
+        return item
+    def __getitem__(self, idx):
+        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
+        item["labels"] = torch.tensor(self.labels[idx], dtype=torch.float)
+        return item
+ 
+
+# Label map used across modules
+label_map = {
+    0.0: "no",
+    0.5: "plausibly",
+    1.0: "yes"
+}
+
+# Function to map probabilities to 3 classes
+# (0.0, 0.5, 1.0) based on thresholds
+def map_to_3_classes(prob_array, low, high):
+    """Map probabilities to 0.0, 0.5, 1.0 using thresholds."""
+    mapped = np.zeros_like(prob_array)
+    mapped[(prob_array > low) & (prob_array <= high)] = 0.5
+    mapped[prob_array > high] = 1.0
+    return mapped
+
+def convert_to_label_strings(array):
+    """Convert float label array to list of strings."""
+    return [label_map[val] for val in array.flatten()]
+
+def tune_thresholds(probs, true_labels, verbose=True):
+    """Search for best (low, high) thresholds by macro F1 score."""
+    best_macro_f1 = 0.0
+    best_low, best_high = 0.0, 0.0
+
+    for low in np.arange(0.2, 0.5, 0.05):
+        for high in np.arange(0.55, 0.8, 0.05):
+            if high <= low:
+                continue
+
+            pred_soft = map_to_3_classes(probs, low, high)
+            pred_str = convert_to_label_strings(pred_soft)
+            true_str = convert_to_label_strings(true_labels)
+
+            _, _, f1, _ = precision_recall_fscore_support(
+                true_str, pred_str,
+                labels=["no", "plausibly", "yes"],
+                average="macro",
+                zero_division=0
+            )
+            if verbose:
+                print(f"low={low:.2f}, high={high:.2f} -> macro F1={f1:.3f}")
+            if f1 > best_macro_f1:
+                best_macro_f1 = f1
+                best_low, best_high = low, high
+
+    return best_low, best_high, best_macro_f1
+
+#  Convert label values to soft scores: "yes" = 1.0, "plausibly" = 0.5, others = 0.0
+def label_row_soft(row):
+    labels = []
+    for col in label_columns:
+        val = str(row[col]).strip().lower()
+        if val == "yes":
+            labels.append(1.0)
+        elif val == "plausibly":
+            labels.append(0.5)
+        else:
+            labels.append(0.0)
+    return labels
+