trying with balanced dataset

app.py

@@ -11,6 +11,7 @@ from transformers import AutoTokenizer, AutoModel
 from sklearn.ensemble import RandomForestClassifier
 from sklearn.preprocessing import StandardScaler
 from sklearn.metrics import accuracy_score, f1_score
+from sklearn.model_selection import train_test_split
 
 # =========================
 # 1. DOWNLOAD DATASET
@@ -31,12 +32,12 @@ if not os.path.exists(extract_path):
         zip_ref.extractall(extract_path)
 
 # =========================
-# 2. LOAD BERT MODEL
+# 2. LOAD LIGHTWEIGHT BERT
 # =========================
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-tokenizer = AutoTokenizer.from_pretrained("bert-base-multilingual-cased")
-bert = AutoModel.from_pretrained("bert-base-multilingual-cased").to(device)
+tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+bert = AutoModel.from_pretrained("distilbert-base-uncased").to(device)
 bert.eval()
 
 # =========================
@@ -69,14 +70,29 @@ def get_label(folder):
     return 1 if "_P" in folder else 0
 
 # =========================
-# 4. TRAIN MODEL (LIGHT)
+# 4. BUILD BALANCED DATASET
 # =========================
-print("Training model...")
+print("Building balanced dataset...")
 
-X, y = [], []
+all_folders = [f for f in os.listdir(extract_path) if os.path.isdir(os.path.join(extract_path, f))]
+
+p_folders = [f for f in all_folders if "_P" in f]
+c_folders = [f for f in all_folders if "_C" in f]
+
+# 🔥 pick equal samples
+num_samples = min(10, len(p_folders), len(c_folders))
 
-folders = [f for f in os.listdir(extract_path) if os.path.isdir(os.path.join(extract_path, f))]
-folders = folders[:20]  # 🔥 small subset for speed
+p_folders = p_folders[:num_samples]
+c_folders = c_folders[:num_samples]
+
+folders = p_folders + c_folders
+
+print(f"Using {len(p_folders)} depressed and {len(c_folders)} control samples")
+
+# =========================
+# 5. FEATURE EXTRACTION
+# =========================
+X, y = [], []
 
 for folder in folders:
     path = os.path.join(extract_path, folder)
@@ -93,43 +109,43 @@ for folder in folders:
 X = np.array(X)
 y = np.array(y)
 
+# =========================
+# 6. TRAIN-TEST SPLIT
+# =========================
+X_train, X_test, y_train, y_test = train_test_split(
+    X, y, test_size=0.3, random_state=42, stratify=y
+)
+
 scaler = StandardScaler()
-X = scaler.fit_transform(X)
+X_train = scaler.fit_transform(X_train)
+X_test = scaler.transform(X_test)
+
+# =========================
+# 7. TRAIN MODEL
+# =========================
+print("Training model...")
 
 model = RandomForestClassifier(n_estimators=50, random_state=42)
-model.fit(X, y)
+model.fit(X_train, y_train)
 
 print("Model ready!")
 
 # =========================
-# 5. RUN PREDICTIONS
+# 8. RUN EVALUATION
 # =========================
 def run_on_dataset():
 
-    results = []
-    preds = []
-    labels = []
-
-    for i, folder in enumerate(folders):
-        path = os.path.join(extract_path, folder)
-
-        text = load_text(path)
-        text_feat = get_text_embedding(text)
-        audio_feat = get_audio_features(path)
-
-        x = np.concatenate([text_feat, audio_feat])
-        x = scaler.transform([x])
+    preds = model.predict(X_test)
 
-        pred = model.predict(x)[0]
-        label = get_label(folder)
+    acc = accuracy_score(y_test, preds)
+    f1 = f1_score(y_test, preds)
 
-        preds.append(pred)
-        labels.append(label)
-
-        results.append(f"{folder} → {'Depressed' if pred==1 else 'Control'}")
-
-    acc = accuracy_score(labels, preds)
-    f1 = f1_score(labels, preds)
+    results = []
+    for i, pred in enumerate(preds):
+        label = y_test[i]
+        results.append(
+            f"Sample {i+1} → Pred: {'Depressed' if pred else 'Control'} | True: {'Depressed' if label else 'Control'}"
+        )
 
     output = "\n".join(results)
     output += f"\n\nAccuracy: {acc:.3f}"
@@ -138,14 +154,14 @@ def run_on_dataset():
     return output
 
 # =========================
-# 6. GRADIO UI
+# 9. GRADIO UI
 # =========================
 app = gr.Interface(
     fn=run_on_dataset,
     inputs=[],
     outputs="text",
     title="Multimodal Depression Detection",
-    description="Automatically trains on dataset and evaluates predictions (text + audio)"
+    description="Balanced dataset training + evaluation (text + audio)"
 )
 
 app.launch()