Fix: Handle 'tweet' column in Arabic dataset correctly

app.py

@@ -17,8 +17,12 @@ 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)
+    text_column = 'tweet' if 'tweet' in examples else 'text'
+    
     # Tokenize the Arabic text
-    encoding = tokenizer(examples['text'], truncation=True, padding='max_length', max_length=128)
+    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']
@@ -26,8 +30,8 @@ def preprocess_function(examples):
         encoding['labels'] = examples['sentiment']
     return encoding
 
-# Preprocess the dataset
-tokenized_dataset = dataset.map(preprocess_function, batched=True, remove_columns=['text'])
+# 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)"""
@@ -53,32 +57,32 @@ def train_model(epochs, batch_size, learning_rate):
         # Start training
         trainer.train()
         
-        return "\u270d✅ \u062aم التدريب بنجاح!\n" + \
-                f"النموذج محفوظ في ./results\nمعدل التعلم: {learning_rate}\nعدد الحقب: {epochs}\nBatch Size: {batch_size}"
+        return "\u270d✅ تم التدريب بنجاح!\n" + \
+        f"النموذج محفوظ في ./results\nمعدل التعلم: {learning_rate}\nعدد الحقب: {epochs}\nBatch Size: {batch_size}"
     except Exception as e:
         return f"❌ خطأ أثناء التدريب: {str(e)}"
 
 # Create Gradio interface
 with gr.Blocks(title="DistilBERT Arabic Sentiment Training") as demo:
     gr.Markdown("""
-    # 🚀 تدريب نموذج DistilBERT العربي
-    
-    ضبط نموذج **DistilBERT** على تحليل المشاعر باللغة العربية (اللهجة السعودية)
-    
-    ### معلومات النموذج:
-    - **النموذج الأساسي**: distilbert-base-multilingual-cased (67M معامل)
-    - **المهمة**: تصنيف النصوص (المتعد اللغات)
-    - **قاعدة البيانات**: arbml/Arabic_Sentiment_Twitter_Corpus (58.8k مثال)
-    - **اللغة**: العربية (اللهجة السعودية والخليجية)
-    """)  
+# 🚀 تدريب نموذج DistilBERT العربي
+
+ضبط نموذج **DistilBERT** على تحليل المشاعر باللغة العربية (اللهجة السعودية)
+
+### معلومات النموذج:
+- **النموذج الأساسي**: distilbert-base-multilingual-cased (67M معامل)
+- **المهمة**: تصنيف النصوص (المتعد اللغات)
+- **قاعدة البيانات**: arbml/Arabic_Sentiment_Twitter_Corpus (58.8k مثال)
+- **اللغة**: العربية (اللهجة السعودية والخليجية)
+""")
     
     with gr.Row():
         with gr.Column():
             gr.Markdown("### إعدادات التدريب")
-            epochs = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="\u0639دد الحقب (Epochs)")
+            epochs = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="عدد الحقب (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)
@@ -91,12 +95,12 @@ with gr.Blocks(title="DistilBERT Arabic Sentiment Training") as demo:
     )
     
     gr.Markdown("""
-    ### تفاصيل التدريب:
-    - **مرحلة البناء**: GPU مجاني (مباشر عبر Hugging Face Spaces)
-    - **وقت المتوقع**: 5-10 دقائق (GPU) أو 15-20 دقيقة (CPU)
-    - **مخرجات النموذج**: محفوظ عند ./results
-    - **الاستخدام**: النصوص العربية فقط
-    """)
+### تفاصيل التدريب:
+- **مرحلة البناء**: GPU مجاني (مباشر عبر Hugging Face Spaces)
+- **وقت المتوقع**: 5-10 دقائق (GPU) أو 15-20 دقيقة (CPU)
+- **مخرجات النموذج**: محفوظ عند ./results
+- **الاستخدام**: النصوص العربية فقط
+""")
 
 if __name__ == "__main__":
     demo.launch()