Update train_utils.py

train_utils.py

@@ -287,24 +287,44 @@ def predict_probabilities(model, loader):
     all_probabilities = [[] for _ in range(len(LABEL_COLUMNS))]
 
     with torch.no_grad():
-        for batch in tqdm(loader, desc="Predicting Probabilities"):
-            # Unpack batch, ignoring labels as we only need inputs
-            if len(batch) == 2:
-                inputs, _ = batch
-                input_ids = inputs['input_ids'].to(DEVICE)
-                attention_mask = inputs['attention_mask'].to(DEVICE)
-                outputs = model(input_ids, attention_mask)
-            elif len(batch) == 3:
-                inputs, metadata, _ = batch
-                input_ids = inputs['input_ids'].to(DEVICE)
-                attention_mask = inputs['attention_mask'].to(DEVICE)
-                metadata = metadata.to(DEVICE)
-                outputs = model(input_ids, attention_mask, metadata)
-            else:
-                raise ValueError("Unsupported batch format.")
+        # Check if we're in a non-interactive environment (like a server)
+        import sys
+        is_interactive = hasattr(sys, 'ps1') or sys.stdout.isatty()
+        
+        if is_interactive:
+            # Use tqdm with proper error handling
+            try:
+                for batch in tqdm(loader, desc="Predicting Probabilities", leave=False):
+                    _process_batch_for_prediction(batch, model, all_probabilities)
+            except Exception as e:
+                print(f"Warning: Error during prediction with progress bar: {str(e)}")
+                # Fallback: process without progress bar
+                for batch in loader:
+                    _process_batch_for_prediction(batch, model, all_probabilities)
+        else:
+            # Non-interactive environment, process without progress bar
+            for batch in loader:
+                _process_batch_for_prediction(batch, model, all_probabilities)
+    
+    return all_probabilities
+
+def _process_batch_for_prediction(batch, model, all_probabilities):
+    """Helper function to process a batch for prediction"""
+    if len(batch) == 2:
+        inputs, _ = batch
+        input_ids = inputs['input_ids'].to(DEVICE)
+        attention_mask = inputs['attention_mask'].to(DEVICE)
+        outputs = model(input_ids, attention_mask)
+    elif len(batch) == 3:
+        inputs, metadata, _ = batch
+        input_ids = inputs['input_ids'].to(DEVICE)
+        attention_mask = inputs['attention_mask'].to(DEVICE)
+        metadata = metadata.to(DEVICE)
+        outputs = model(input_ids, attention_mask, metadata)
+    else:
+        raise ValueError("Unsupported batch format.")
 
-            for i, out_logits in enumerate(outputs):
-                # Apply softmax to logits to get probabilities
-                probs = torch.softmax(out_logits, dim=1).cpu().numpy()
-                all_probabilities[i].extend(probs)
-    return all_probabilities
+    for i, out_logits in enumerate(outputs):
+        # Apply softmax to logits to get probabilities
+        probs = torch.softmax(out_logits, dim=1).cpu().numpy()
+        all_probabilities[i].extend(probs)