Update app.py

app.py

@@ -40,6 +40,11 @@ def setup_training():
     # Print dataset structure
     logging.info(f"Dataset columns: {dataset['train'].column_names}")
 
+    # Determine the number of unique labels
+    unique_labels = dataset['train'].unique(config['target_column'])
+    num_labels = len(unique_labels)
+    logging.info(f"Number of unique labels: {num_labels}")
+
     logging.info(f"Dataset loaded. Train size: {len(dataset['train'])}, Test size: {len(dataset['test'])}")
 
     # Load tokenizer and model
@@ -47,7 +52,7 @@ def setup_training():
     tokenizer = AutoTokenizer.from_pretrained(config['model_name'])
     model = AutoModelForSequenceClassification.from_pretrained(
         config['model_name'], 
-        num_labels=len(dataset['train'].features[config['target_column']].names)
+        num_labels=num_labels
     )
 
     # Tokenize the dataset