Optimize: Lazy load dataset and model for faster app startup

app.py

@@ -2,40 +2,51 @@ import gradio as gr
 import torch
 from datasets import load_dataset
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
-import numpy as np
 
-# Load the Arabic sentiment dataset (Saudi dialect from Twitter)
-try:
-    dataset = load_dataset('arbml/Arabic_Sentiment_Twitter_Corpus')
-    print(f"Dataset loaded with {len(dataset['train'])} training examples")
-except:
-    print("Loading alternative Arabic dataset...")
-    dataset = load_dataset('asas-ai/Arabic_Sentiment_Twitter_Corpus')
+# Global variables for caching
+dataset = None
+tokenizer = None
+model = None
+tokenized_dataset = None
 
-# Load tokenizer and model (supports Arabic)
-tokenizer = AutoTokenizer.from_pretrained('distilbert-base-multilingual-cased')
-model = AutoModelForSequenceClassification.from_pretrained('distilbert-base-multilingual-cased', num_labels=3)
+def load_resources():
+    """Load dataset, tokenizer, and model on demand"""
+    global dataset, tokenizer, model, tokenized_dataset
+    
+    if dataset is not None:
+        return
+    
+    # Load the Arabic sentiment dataset
+    try:
+        dataset = load_dataset('arbml/Arabic_Sentiment_Twitter_Corpus')
+    except:
+        dataset = load_dataset('asas-ai/Arabic_Sentiment_Twitter_Corpus')
+    
+    # Load tokenizer and model
+    tokenizer = AutoTokenizer.from_pretrained('distilbert-base-multilingual-cased')
+    model = AutoModelForSequenceClassification.from_pretrained('distilbert-base-multilingual-cased', num_labels=3)
 
 def preprocess_function(examples):
-    # Check which column contains the text (tweet or text)
+    """Tokenize and preprocess examples"""
     text_column = 'tweet' if 'tweet' in examples else 'text'
-    
-    # Tokenize the Arabic text
     encoding = tokenizer(examples[text_column], truncation=True, padding='max_length', max_length=128)
     
-    # Map label to indices
     if 'label' in examples:
         encoding['labels'] = examples['label']
     elif 'sentiment' in examples:
         encoding['labels'] = examples['sentiment']
     return encoding
 
-# Preprocess the dataset - only keep label and input_ids columns
-tokenized_dataset = dataset.map(preprocess_function, batched=True, remove_columns=dataset['train'].column_names)
-
 def train_model(epochs, batch_size, learning_rate):
     """Fine-tune DistilBERT on Arabic sentiment dataset (Saudi dialect)"""
     try:
+        load_resources()
+        
+        # Preprocess dataset if not already done
+        global tokenized_dataset
+        if tokenized_dataset is None:
+            tokenized_dataset = dataset.map(preprocess_function, batched=True, remove_columns=dataset['train'].column_names)
+        
         training_args = TrainingArguments(
             output_dir='./results',
             num_train_epochs=int(epochs),
@@ -57,37 +68,36 @@ def train_model(epochs, batch_size, learning_rate):
         # Start training
         trainer.train()
         
-        return "\u270d✅ تم التدريب بنجاح!\n" + \
-        f"النموذج محفوظ في ./results\nمعدل التعلم: {learning_rate}\nعدد الحقب: {epochs}\nBatch Size: {batch_size}"
+        return f"\u270d\u2705 \u062a\u0645 \u0627\u0644\u062a\u062f\u0631\u064a\u0628 \u0628\u0646\u062c\u0627\u062d!\n\u0627\u0644\u0646\u0645\u0648\u0630\u062c \u0645\u062d\u0641\u0648\u0638 \u0641\u064a ./results\n\u0645\u0639\u062f\u0644 \u0627\u0644\u062a\u0639\u0644\u0645: {learning_rate}\n\u0639\u062f\u062f \u0627\u0644\u062d\u0642\u0628: {epochs}\nBatch Size: {batch_size}"
     except Exception as e:
-        return f"❌ خطأ أثناء التدريب: {str(e)}"
+        return f"\u274c \u062e\u0637\u0623 \u0623\u062b\u0646\u0627\u0621 \u0627\u0644\u062a\u062f\u0631\u064a\u0628: {str(e)}"
 
 # Create Gradio interface
 with gr.Blocks(title="DistilBERT Arabic Sentiment Training") as demo:
     gr.Markdown("""
-# 🚀 تدريب نموذج DistilBERT العربي
+# \ud83d\ude80 \u062a\u062f\u0631\u064a\u0628 \u0646\u0645\u0648\u0630\u062c DistilBERT \u0627\u0644\u0639\u0631\u0628\u064a
 
-ضبط نموذج **DistilBERT** على تحليل المشاعر باللغة العربية (اللهجة السعودية)
+\u0636\u0628\u0637 \u0646\u0645\u0648\u0630\u062c **DistilBERT** \u0639\u0644\u0649 \u062a\u062d\u0644\u064a\u0644 \u0627\u0644\u0645\u0634\u0627\u0639\u0631 \u0628\u0627\u0644\u0644\u063a\u0629 \u0627\u0644\u0639\u0631\u0628\u064a\u0629 (\u0627\u0644\u0644\u0647\u062c\u0629 \u0627\u0644\u0633\u0639\u0648\u062f\u064a\u0629)
 
-### معلومات النموذج:
-- **النموذج الأساسي**: distilbert-base-multilingual-cased (67M معامل)
-- **المهمة**: تصنيف النصوص (المتعد اللغات)
-- **قاعدة البيانات**: arbml/Arabic_Sentiment_Twitter_Corpus (58.8k مثال)
-- **اللغة**: العربية (اللهجة السعودية والخليجية)
+### \u0645\u0639\u0644\u0648\u0645\u0627\u062a \u0627\u0644\u0646\u0645\u0648\u0630\u062c:
+- **\u0627\u0644\u0646\u0645\u0648\u0630\u062c \u0627\u0644\u0623\u0633\u0627\u0633\u064a**: distilbert-base-multilingual-cased (67M \u0645\u0639\u0627\u0645\u0644)
+- **\u0627\u0644\u0645\u0647\u0645\u0629**: \u062a\u0635\u0646\u064a\u0641 \u0627\u0644\u0646\u0635\u0648\u0635 (\u0627\u0644\u0645\u062a\u0639\u062f \u0627\u0644\u0644\u063a\u0627\u062a)
+- **\u0642\u0627\u0639\u062f\u0629 \u0627\u0644\u0628\u064a\u0627\u0646\u0627\u062a**: arbml/Arabic_Sentiment_Twitter_Corpus (58.8k \u0645\u062b\u0627\u0644)
+- **\u0627\u0644\u0644\u063a\u0629**: \u0627\u0644\u0639\u0631\u0628\u064a\u0629 (\u0627\u0644\u0644\u0647\u062c\u0629 \u0627\u0644\u0633\u0639\u0648\u062f\u064a\u0629 \u0648\u0627\u0644\u062e\u0644\u064a\u062c\u064a\u0629)
 """)
     
     with gr.Row():
         with gr.Column():
-            gr.Markdown("### إعدادات التدريب")
-            epochs = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="عدد الحقب (Epochs)")
+            gr.Markdown("### \u0625\u0639\u062f\u0627\u062f\u0627\u062a \u0627\u0644\u062a\u062f\u0631\u064a\u0628")
+            epochs = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="\u0639\u062f\u062f \u0627\u0644\u062d\u0642\u0628 (Epochs)")
             batch_size = gr.Slider(minimum=8, maximum=64, value=32, step=8, label="Batch Size")
             learning_rate = gr.Slider(minimum=1e-5, maximum=1e-3, value=2e-5, step=1e-5, label="Learning Rate")
         
         with gr.Column():
-            gr.Markdown("### حالة التدريب")
-            output_text = gr.Textbox(label="المخرجات", lines=10, interactive=False)
+            gr.Markdown("### \u062d\u0627\u0644\u0629 \u0627\u0644\u062a\u062f\u0631\u064a\u0628")
+            output_text = gr.Textbox(label="\u0627\u0644\u0645\u062e\u0631\u062c\u0627\u062a", lines=10, interactive=False)
     
-    train_button = gr.Button("🔥 بدء التدريب", variant="primary", scale=2)
+    train_button = gr.Button("\ud83d\udd25 \u0628\u062f\u0621 \u0627\u0644\u062a\u062f\u0631\u064a\u0628", variant="primary", scale=2)
     train_button.click(
         fn=train_model,
         inputs=[epochs, batch_size, learning_rate],
@@ -95,11 +105,12 @@ with gr.Blocks(title="DistilBERT Arabic Sentiment Training") as demo:
     )
     
     gr.Markdown("""
-### تفاصيل التدريب:
-- **مرحلة البناء**: GPU مجاني (مباشر عبر Hugging Face Spaces)
-- **وقت المتوقع**: 5-10 دقائق (GPU) أو 15-20 دقيقة (CPU)
-- **مخرجات النموذج**: محفوظ عند ./results
-- **الاستخدام**: النصوص العربية فقط
+### \u062a\u0641\u0627\u0635\u064a\u0644 \u0627\u0644\u062a\u062f\u0631\u064a\u0628:
+- **\u0645\u0631\u062d\u0644\u0629 \u0627\u0644\u0628\u0646\u0627\u0621**: GPU \u0645\u062c\u0627\u0646\u064a (\u0645\u0628\u0627\u0634\u0631 \u0639\u0628\u0631 Hugging Face Spaces)
+- **\u0648\u0642\u062a \u0627\u0644\u062a\u062d\u0645\u064a\u0644**: 5-10 \u062f\u0642\u0627\u0626\u0642 (GPU) \u0623\u0648 15-20 \u062f\u0642\u064a\u0642\u0629 (CPU)
+- **\u0648\u0642\u062a \u0627\u0644\u062a\u062f\u0631\u064a\u0628**: \u064a\u0639\u062a\u0645\u062f \u0639\u0644\u0649 \u0639\u062f\u062f \u0627\u0644\u062d\u0642\u0628 \u0648Batch Size
+- **\u0645\u062e\u0631\u062c\u0627\u062a \u0627\u0644\u0646\u0645\u0648\u0630\u062c**: \u0645\u062d\u0641\u0648\u0638 \u0639\u0646\u062f ./results
+- **\u0627\u0644\u0627\u0633\u062a\u062e\u062f\u0627\u0645**: \u0627\u0644\u0646\u0635\u0648\u0635 \u0627\u0644\u0639\u0631\u0628\u064a\u0629 \u0641\u0642\u0637
 """)
 
 if __name__ == "__main__":