feat: persist model and logs in Hugging Face Space + add model push to hub

- Updated paths to save model and logs to /home/user/app (persistent in Spaces)
- Modified logging to stream to both file and UI log buffer
- Updated model saving/loading to use MODEL_DIR inside the persistent path
- Added push_model_to_hub() to upload trained model/tokenizer to Hugging Face Hub
- Extended Gradio UI with 'Evaluate Model' and 'Push Model to Hub' buttons

app.py

@@ -1,17 +1,20 @@
 import gradio as gr
-from train_abuse_model import run_training
+from train_abuse_model import run_training, evaluate_saved_model, push_model_to_hub
 
 with gr.Blocks() as demo:
     gr.Markdown("## 🧠 Abuse Detection Fine-Tuning App")
-    gr.Markdown(
-        "⚠️ **Important:** Keep this tab open and prevent your computer from sleeping while training runs."
-    )
+    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")
 
-    output_box = gr.Textbox(label="Live Training Logs", lines=25, interactive=False)
+    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)
 
 if __name__ == "__main__":
     demo.launch()

train_abuse_model.py

@@ -5,6 +5,7 @@ import logging
 import io
 import os
 import gradio as gr  # ✅ required for progress bar
+from pathlib import Path
 
 # Python standard + ML packages
 import pandas as pd
@@ -30,18 +31,23 @@ from transformers import (
     TrainingArguments
 )
 
+PERSIST_DIR = Path("/home/user/app")
+MODEL_DIR = PERSIST_DIR / "saved_model"
+LOG_FILE = PERSIST_DIR / "training.log"
+
 # configure logging 
-log_buffer = io.StringIO()  
+log_buffer = io.StringIO()
 logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s - %(levelname)s - %(message)s",
     handlers=[
-        logging.FileHandler("training.log"),       # to file
-        logging.StreamHandler(log_buffer)          # to in-memory buffer
+        logging.FileHandler(LOG_FILE),
+        logging.StreamHandler(log_buffer)
     ]
 )
 logger = logging.getLogger(__name__)
 
+
 # Check versions
 logger.info("Transformers version:", transformers.__version__)
 
@@ -50,6 +56,9 @@ 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):
@@ -127,33 +136,81 @@ def tune_thresholds(probs, true_labels, verbose=True):
 def evaluate_model_with_thresholds(trainer, test_dataset):
     """Run full evaluation with automatic threshold tuning."""
     logger.info("\n🔍 Running model predictions...")
+    yield "\n🔍 Running model predictions..."
+
     predictions = trainer.predict(test_dataset)
     probs = torch.sigmoid(torch.tensor(predictions.predictions)).numpy()
     true_soft = np.array(predictions.label_ids)
 
     logger.info("\n🔎 Tuning thresholds...")
+    yield "\n🔎 Tuning thresholds..."
     best_low, best_high, best_f1 = tune_thresholds(probs, true_soft)
 
     logger.info(f"\n✅ Best thresholds: low={best_low:.2f}, high={best_high:.2f} (macro F1={best_f1:.3f})")
+    yield f"\n✅ Best thresholds: low={best_low:.2f}, high={best_high:.2f} (macro F1={best_f1:.3f})"    
 
     final_pred_soft = map_to_3_classes(probs, best_low, best_high)
     final_pred_str = convert_to_label_strings(final_pred_soft)
     true_str = convert_to_label_strings(true_soft)
 
     logger.info("\n📊 Final Evaluation Report (multi-class per label):\n")
+    yield "\n📊 Final Evaluation Report (multi-class per label):\n "
     logger.info(classification_report(
         true_str,
         final_pred_str,
         labels=["no", "plausibly", "yes"],
+        digits=3,
         zero_division=0
     ))
+    yield classification_report(
+        true_str,
+        final_pred_str,
+        labels=["no", "plausibly", "yes"],
+        digits=3,
+        zero_division=0
+    )
+def load_saved_model_and_tokenizer():
+    tokenizer = DebertaV2Tokenizer.from_pretrained(MODEL_DIR)
+    model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR).to(device)
+    return tokenizer, model
+
+def evaluate_saved_model(progress=gr.Progress(track_tqdm=True)):
+    if os.path.exists("saved_model/"):
+        yield "✅ Trained model found! Skipping training...\n"
+    else:
+        yield "❌ No trained model found. Please train the model first.\n"
+        return
+    try:
+        logger.info("🔍 Loading saved model for evaluation...")
+        yield "🔍 Loading saved model for evaluation...\n"
+
+        tokenizer, model = load_saved_model_and_tokenizer()
+        test_dataset = AbuseDataset(test_texts, test_labels, tokenizer)
+
+        trainer = Trainer(
+            model=model,
+            args=TrainingArguments(
+                output_dir="./results_eval",
+                per_device_eval_batch_size=4,
+                logging_dir="./logs_eval",
+                disable_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
+
+        logger.info("✅ Evaluation complete.\n")
+        yield "\n✅ Evaluation complete.\n"
+
+    except Exception as e:
+        logger.exception(f"❌ Evaluation failed: {e}")
+        yield f"❌ Evaluation failed: {e}\n"
 
-    return {
-        "thresholds": (best_low, best_high),
-        "macro_f1": best_f1,
-        "true_labels": true_str,
-        "pred_labels": final_pred_str
-    }
 
 token = os.environ.get("HF_TOKEN")  # Reads my token from a secure hf secret
 
@@ -202,6 +259,10 @@ train_texts, val_texts, train_labels, val_labels = train_test_split(
 model_name = "microsoft/deberta-v3-base"
 
 def run_training(progress=gr.Progress(track_tqdm=True)):
+    if os.path.exists("saved_model/"):
+        yield "✅ Trained model found! Skipping training...\n"
+        yield evaluate_saved_model()
+        return
     yield "🚀 Starting training...\n"
     try:
         logger.info("Starting training run...")
@@ -269,11 +330,10 @@ def run_training(progress=gr.Progress(track_tqdm=True)):
         trainer.train()
 
         # Save the model and tokenizer
-        if not os.path.exists("saved_model/"):
-            os.makedirs("saved_model/")
-        model.save_pretrained("saved_model/")
-        tokenizer.save_pretrained("saved_model/")
-        
+        MODEL_DIR.mkdir(parents=True, exist_ok=True)
+        model.save_pretrained(MODEL_DIR)
+        tokenizer.save_pretrained(MODEL_DIR)
+
         logger.info(" Training completed and model saved.")
         yield "🎉 Training complete! Model saved.\n"
 
@@ -284,7 +344,6 @@ def run_training(progress=gr.Progress(track_tqdm=True)):
     # Evaluation
     try:
         if 'trainer' in locals():
-            label_map = {0.0: "no", 0.5: "plausibly", 1.0: "yes"}
             evaluate_model_with_thresholds(trainer, test_dataset)
             logger.info("Evaluation completed")
     except Exception as e:
@@ -292,3 +351,15 @@ def run_training(progress=gr.Progress(track_tqdm=True)):
     log_buffer.seek(0)
     return log_buffer.read()
 
+def push_model_to_hub():
+    try:
+        logger.info("🔄 Pushing model to Hugging Face Hub...")
+        tokenizer, model = load_saved_model_and_tokenizer()
+        model.push_to_hub("rshakked/safe-talk", use_auth_token=token)
+        tokenizer.push_to_hub("rshakked/safe-talk", use_auth_token=token)
+        return "✅ Model pushed to hub successfully!"
+    except Exception as e:
+        logger.exception("❌ Failed to push model to hub.")
+        return f"❌ Failed to push model: {e}"
+
+