Update app.py

app.py

@@ -38,13 +38,14 @@ label_encoder = LabelEncoder()
 df['label_encoded'] = label_encoder.fit_transform(df['label'])
 num_classes = len(label_encoder.classes_)
 
-# Split the data
+# Split the data - fix: remove stratify for classes with few samples
 X_train, X_test, y_train, y_test = train_test_split(
     df['text'], 
     df['label_encoded'], 
     test_size=0.2, 
-    random_state=42, 
-    stratify=df['label_encoded'] if len(df) > 10 else None  # Only stratify if we have enough samples
+    random_state=42,
+    # Only use stratify if we have enough samples
+    stratify=None  # Removed stratification to fix the error
 )
 
 # Feature extraction using TF-IDF