Spaces:
Sleeping
Sleeping
Update tasks/text.py
Browse files- tasks/text.py +55 -31
tasks/text.py
CHANGED
@@ -23,35 +23,55 @@ ROUTE = "/text"
|
|
23 |
|
24 |
class TextClassifier:
|
25 |
def __init__(self):
|
26 |
-
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
32 |
-
|
33 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
34 |
|
35 |
def process_batch(self, batch: List[str], batch_idx: int) -> Tuple[List[int], int]:
|
36 |
-
"""
|
37 |
-
|
38 |
-
|
39 |
-
|
40 |
-
|
41 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
42 |
|
43 |
-
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
-
|
48 |
-
|
49 |
-
|
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 |
-
|
137 |
-
|
|
|
138 |
except Exception as e:
|
139 |
print(f"Failed to get results for batch {batch_idx}: {e}")
|
140 |
-
|
|
|
141 |
|
142 |
# Flatten predictions while maintaining order
|
143 |
-
predictions = [
|
144 |
-
|
|
|
|
|
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
|