debugging prediction request and response

model.py

@@ -105,64 +105,59 @@ class BertClassifier(LabelStudioMLBase):
         return self
 
     def predict(self, tasks, **kwargs):
-        """Generate predictions for a list of tasks."""
+        """
+        Tasks is a list of tasks with the following fields:
+        {
+            "id": 123,
+            "data": {
+                "text": "Example text"
+            }
+        }
+        """
         logger.info("=== PREDICT METHOD CALLED ===")
+        logger.info(f"Received tasks: {json.dumps(tasks, indent=2)}")
         logger.info(f"Number of tasks: {len(tasks)}")
-        
-        # Verify model state
-        model_path = os.path.join(self.model_dir, 'model_state.pt')
-        if os.path.exists(model_path):
-            logger.info(f"✓ Using trained model from: {model_path}")
-        else:
-            logger.info("✗ No trained model found, using initial state")
-        
+
         predictions = []
-        for task in tasks:
-            task_id = task['id']
-            text = task['data']['text']
-            logger.info(f"Processing task {task_id} - Text: {text[:50]}...")
-            
+        
+        for task_index, task in enumerate(tasks, 1):
             try:
-                # Prepare the text
-                inputs = self.tokenizer(
-                    text,
-                    truncation=True,
-                    padding=True,
-                    return_tensors='pt'
-                ).to(self.device)
+                # Log the specific task being processed
+                logger.info(f"Processing task {task_index} - Text: {task['data'].get('text', '')[:20]}...")
                 
-                # Get model predictions
-                self._model.eval()
-                with torch.no_grad():
-                    outputs = self._model(**inputs)
-                    
-                # Get predicted category and confidence
-                probabilities = torch.softmax(outputs.logits, dim=1)
-                confidence, predicted_idx = torch.max(probabilities, dim=1)
-                predicted_category = self.categories[predicted_idx.item()]
-                confidence = confidence.item()
+                # Log model state
+                model_path = os.path.join(self.model_dir, 'model_state.pt')
+                if os.path.exists(model_path):
+                    logger.info("✓ Using trained model")
+                else:
+                    logger.info("✗ No trained model found, using initial state")
                 
-                logger.info(f"Predicted category: {predicted_category} with confidence: {confidence:.4f}")
+                # Get model prediction
+                predicted_label, confidence = self._get_prediction(task['data']['text'])
+                logger.info(f"Predicted category: {predicted_label} with confidence: {confidence:.4f}")
                 
-                # Format prediction for Label Studio
-                predictions.append({
+                # Format the prediction for Label Studio
+                prediction = {
                     'result': [{
                         'from_name': 'sentiment',
                         'to_name': 'text',
                         'type': 'choices',
                         'value': {
-                            'choices': [predicted_category]
-                        }
+                            'choices': [predicted_label]
+                        },
+                        'score': confidence
                     }],
-                    'score': confidence,
-                    'model_version': 'bert-base-uncased-v1'
-                })
+                    'model_version': self.model_version,
+                    'task': task['id']
+                }
+                predictions.append(prediction)
                 
             except Exception as e:
-                logger.error(f"Error predicting task {task_id}: {str(e)}")
+                logger.error(f"Error predicting task {task_index}: {str(e)}")
                 continue
-                
+        
         logger.info(f"Returning {len(predictions)} predictions")
+        logger.info(f"Predictions: {json.dumps(predictions, indent=2)}")
         return predictions
 
     def fit(self, event_data, data=None, **kwargs):