camillebrl commited on
Commit
c01695c
·
verified ·
1 Parent(s): 92fa037

Update tasks/text.py

Browse files
Files changed (1) hide show
  1. tasks/text.py +55 -31
tasks/text.py CHANGED
@@ -23,35 +23,55 @@ ROUTE = "/text"
23
 
24
  class TextClassifier:
25
  def __init__(self):
26
- self.config = AutoConfig.from_pretrained("camillebrl/ModernBERT-envclaims-overfit")
27
- self.label2id = self.config.label2id
28
- self.classifier = pipeline(
29
- "text-classification",
30
- "camillebrl/ModernBERT-envclaims-overfit",
31
- device="cpu",
32
- batch_size=16
33
- )
 
 
 
 
 
 
 
 
 
 
 
34
 
35
  def process_batch(self, batch: List[str], batch_idx: int) -> Tuple[List[int], int]:
36
- """
37
- Process a batch of texts and return their predictions along with batch index
38
-
39
- Args:
40
- batch: List of texts to process
41
- batch_idx: Index of the current batch
 
 
 
 
 
 
 
 
 
 
 
 
 
 
42
 
43
- Returns:
44
- Tuple containing list of predictions and batch index
45
- """
46
- try:
47
- print(f"Processing batch {batch_idx} with {len(batch)} items")
48
- batch_preds = self.classifier(list(batch))
49
- predictions = [self.label2id[pred[0]["label"]] for pred in batch_preds]
50
- print(f"Completed batch {batch_idx} with {len(predictions)} predictions")
51
- return predictions, batch_idx
52
- except Exception as e:
53
- print(f"Error in batch {batch_idx}: {str(e)}")
54
- return [], batch_idx
55
 
56
  @router.post(ROUTE, tags=["Text Task"],
57
  description=DESCRIPTION)
@@ -133,15 +153,19 @@ async def evaluate_text(request: TextEvaluationRequest):
133
  batch_idx = future_to_batch[future]
134
  try:
135
  predictions, idx = future.result()
136
- batch_results[idx] = predictions
137
- print(f"Stored results for batch {idx}")
 
138
  except Exception as e:
139
  print(f"Failed to get results for batch {batch_idx}: {e}")
140
- batch_results[batch_idx] = []
 
141
 
142
  # Flatten predictions while maintaining order
143
- predictions = [pred for batch_preds in batch_results for pred in batch_preds]
144
- print(f"Total predictions collected: {len(predictions)}")
 
 
145
 
146
  #--------------------------------------------------------------------------------------------
147
  # YOUR MODEL INFERENCE STOPS HERE
 
23
 
24
  class TextClassifier:
25
  def __init__(self):
26
+ # Add retry mechanism for model initialization
27
+ max_retries = 3
28
+ for attempt in range(max_retries):
29
+ try:
30
+ self.config = AutoConfig.from_pretrained("camillebrl/ModernBERT-envclaims-overfit")
31
+ self.label2id = self.config.label2id
32
+ self.classifier = pipeline(
33
+ "text-classification",
34
+ "camillebrl/ModernBERT-envclaims-overfit",
35
+ device="cpu",
36
+ batch_size=16
37
+ )
38
+ print("Model initialized successfully")
39
+ break
40
+ except Exception as e:
41
+ if attempt == max_retries - 1:
42
+ raise Exception(f"Failed to initialize model after {max_retries} attempts: {str(e)}")
43
+ print(f"Attempt {attempt + 1} failed, retrying...")
44
+ time.sleep(1)
45
 
46
  def process_batch(self, batch: List[str], batch_idx: int) -> Tuple[List[int], int]:
47
+ """Process a batch of texts and return their predictions"""
48
+ max_retries = 3
49
+ for attempt in range(max_retries):
50
+ try:
51
+ print(f"Processing batch {batch_idx} with {len(batch)} items (attempt {attempt + 1})")
52
+ # Process texts one by one in case of errors
53
+ predictions = []
54
+ for text in batch:
55
+ try:
56
+ pred = self.classifier(text)
57
+ pred_label = self.label2id[pred[0]["label"]]
58
+ predictions.append(pred_label)
59
+ except Exception as e:
60
+ print(f"Error processing text in batch {batch_idx}: {str(e)}")
61
+
62
+ if not predictions:
63
+ raise Exception("No predictions generated for batch")
64
+
65
+ print(f"Completed batch {batch_idx} with {len(predictions)} predictions")
66
+ return predictions, batch_idx
67
 
68
+ except Exception as e:
69
+ if attempt == max_retries - 1:
70
+ print(f"Final error in batch {batch_idx}: {str(e)}")
71
+ return [0] * len(batch), batch_idx # Return default predictions instead of empty list
72
+ print(f"Error in batch {batch_idx} (attempt {attempt + 1}): {str(e)}")
73
+ time.sleep(1)
74
+
 
 
 
 
 
75
 
76
  @router.post(ROUTE, tags=["Text Task"],
77
  description=DESCRIPTION)
 
153
  batch_idx = future_to_batch[future]
154
  try:
155
  predictions, idx = future.result()
156
+ if predictions: # Only store non-empty predictions
157
+ batch_results[idx] = predictions
158
+ print(f"Stored results for batch {idx} ({len(predictions)} predictions)")
159
  except Exception as e:
160
  print(f"Failed to get results for batch {batch_idx}: {e}")
161
+ # Use default predictions instead of empty list
162
+ batch_results[batch_idx] = [0] * len(batches[batch_idx])
163
 
164
  # Flatten predictions while maintaining order
165
+ predictions = []
166
+ for batch_preds in batch_results:
167
+ if batch_preds is not None:
168
+ predictions.extend(batch_preds)
169
 
170
  #--------------------------------------------------------------------------------------------
171
  # YOUR MODEL INFERENCE STOPS HERE