Update tasks/text.py

tasks/text.py

@@ -23,35 +23,55 @@ ROUTE = "/text"
 
 class TextClassifier:
     def __init__(self):
-        self.config = AutoConfig.from_pretrained("camillebrl/ModernBERT-envclaims-overfit")
-        self.label2id = self.config.label2id
-        self.classifier = pipeline(
-            "text-classification",
-            "camillebrl/ModernBERT-envclaims-overfit",
-            device="cpu",
-            batch_size=16
-        )
+        # Add retry mechanism for model initialization
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                self.config = AutoConfig.from_pretrained("camillebrl/ModernBERT-envclaims-overfit")
+                self.label2id = self.config.label2id
+                self.classifier = pipeline(
+                    "text-classification",
+                    "camillebrl/ModernBERT-envclaims-overfit",
+                    device="cpu",
+                    batch_size=16
+                )
+                print("Model initialized successfully")
+                break
+            except Exception as e:
+                if attempt == max_retries - 1:
+                    raise Exception(f"Failed to initialize model after {max_retries} attempts: {str(e)}")
+                print(f"Attempt {attempt + 1} failed, retrying...")
+                time.sleep(1)
 
     def process_batch(self, batch: List[str], batch_idx: int) -> Tuple[List[int], int]:
-        """
-        Process a batch of texts and return their predictions along with batch index
-        
-        Args:
-            batch: List of texts to process
-            batch_idx: Index of the current batch
+        """Process a batch of texts and return their predictions"""
+        max_retries = 3
+        for attempt in range(max_retries):
+            try:
+                print(f"Processing batch {batch_idx} with {len(batch)} items (attempt {attempt + 1})")
+                # Process texts one by one in case of errors
+                predictions = []
+                for text in batch:
+                    try:
+                        pred = self.classifier(text)
+                        pred_label = self.label2id[pred[0]["label"]]
+                        predictions.append(pred_label)
+                    except Exception as e:
+                        print(f"Error processing text in batch {batch_idx}: {str(e)}")
+                
+                if not predictions:
+                    raise Exception("No predictions generated for batch")
+                
+                print(f"Completed batch {batch_idx} with {len(predictions)} predictions")
+                return predictions, batch_idx
             
-        Returns:
-            Tuple containing list of predictions and batch index
-        """
-        try:
-            print(f"Processing batch {batch_idx} with {len(batch)} items")
-            batch_preds = self.classifier(list(batch))
-            predictions = [self.label2id[pred[0]["label"]] for pred in batch_preds]
-            print(f"Completed batch {batch_idx} with {len(predictions)} predictions")
-            return predictions, batch_idx
-        except Exception as e:
-            print(f"Error in batch {batch_idx}: {str(e)}")
-            return [], batch_idx
+            except Exception as e:
+                if attempt == max_retries - 1:
+                    print(f"Final error in batch {batch_idx}: {str(e)}")
+                    return [0] * len(batch), batch_idx  # Return default predictions instead of empty list
+                print(f"Error in batch {batch_idx} (attempt {attempt + 1}): {str(e)}")
+                time.sleep(1)
+
 
 @router.post(ROUTE, tags=["Text Task"], 
              description=DESCRIPTION)
@@ -133,15 +153,19 @@ async def evaluate_text(request: TextEvaluationRequest):
             batch_idx = future_to_batch[future]
             try:
                 predictions, idx = future.result()
-                batch_results[idx] = predictions
-                print(f"Stored results for batch {idx}")
+                if predictions:  # Only store non-empty predictions
+                    batch_results[idx] = predictions
+                    print(f"Stored results for batch {idx} ({len(predictions)} predictions)")
             except Exception as e:
                 print(f"Failed to get results for batch {batch_idx}: {e}")
-                batch_results[batch_idx] = []
+                # Use default predictions instead of empty list
+                batch_results[batch_idx] = [0] * len(batches[batch_idx])
 
     # Flatten predictions while maintaining order
-    predictions = [pred for batch_preds in batch_results for pred in batch_preds]
-    print(f"Total predictions collected: {len(predictions)}")
+    predictions = []
+    for batch_preds in batch_results:
+        if batch_preds is not None:
+            predictions.extend(batch_preds)
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE