logs cleanup

model.py

@@ -76,81 +76,42 @@ class BertClassifier(LabelStudioMLBase):
             logger.error("Full error details:", exc_info=True)
 
     def predict(self, tasks, **kwargs):
-        """Make predictions for tasks"""
-        predictions = []
+        """Generate predictions for a list of tasks."""
+        logger.info("=== PREDICT METHOD CALLED ===")
+        logger.info(f"Number of tasks: {len(tasks)}")
         
-        try:
-            logger.info("=== PREDICT METHOD CALLED ===")
-            logger.info(f"Number of tasks received: {len(tasks)}")
-            
-            if self._model is None or self.tokenizer is None:
-                logger.error("Model or tokenizer not initialized")
-                return []
+        # 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 in tasks:
-                logger.info(f"Processing task ID: {task.get('id')}")
-                text = task['data'].get('text', '')
-                logger.info(f"Text to predict: {text[:100]}...")
+            try:
+                # ... prediction code ...
+                logger.info(f"Predicted category: {predicted_category} with confidence: {confidence:.4f}")
+                
+            except Exception as e:
+                logger.error(f"Error predicting task {task_id}: {str(e)}")
+                continue
                 
-                try:
-                    inputs = self.tokenizer(
-                        text,
-                        truncation=True,
-                        padding=True,
-                        return_tensors='pt'
-                    ).to(self.device)
-                    
-                    self._model.eval()
-                    with torch.no_grad():
-                        outputs = self._model(**inputs)
-                        probs = torch.softmax(outputs.logits, dim=1)
-                        predicted_idx = torch.argmax(probs, dim=1).item()
-                        confidence = probs[0][predicted_idx].item()
-                    
-                    predicted_category = self.categories[predicted_idx]
-                    logger.info(f"Predicted category: {predicted_category} with confidence: {confidence:.4f}")
-                    
-                    prediction = {
-                        'result': [{
-                            'from_name': 'sentiment',
-                            'to_name': 'text',
-                            'type': 'choices',
-                            'value': {
-                                'choices': [predicted_category]
-                            }
-                        }],
-                        'score': confidence,
-                        'model_version': 'bert-base-uncased-v1'
-                    }
-                    predictions.append(prediction)
-                    
-                except Exception as e:
-                    logger.error(f"Error processing individual task: {str(e)}")
-                    logger.error("Full error details:", exc_info=True)
-                    predictions.append({
-                        'result': [],
-                        'score': 0,
-                        'model_version': 'bert-base-uncased-v1'
-                    })
-            
-        except Exception as e:
-            logger.error(f"Error in predict: {str(e)}")
-            logger.error("Full error details:", exc_info=True)
-            return []
-        
         logger.info(f"Returning {len(predictions)} predictions")
         return predictions
 
     def fit(self, event_data, data=None, **kwargs):
         """Train the model on a single annotation."""
         start_time = datetime.now()
-        logger.info(f"=== FIT METHOD CALLED ===")
-        logger.info(f"Event data: {event_data}")
-        logger.info(f"Data received: {json.dumps(data, indent=2)}")
+        logger.info("=== FIT METHOD CALLED ===")
+        logger.info(f"Event type: {event_data}")
         
         try:
             if event_data == 'ANNOTATION_CREATED':
-                logger.info("Processing ANNOTATION_CREATED event")
                 annotation = data.get('annotation', {})
                 task = data.get('task', {})
                 
@@ -175,12 +136,12 @@ class BertClassifier(LabelStudioMLBase):
                                 tokenizer=self.tokenizer
                             )
                             train_loader = DataLoader(dataset, batch_size=1)
-                            logger.info("Dataset created successfully")
+                            logger.info("✓ Dataset created")
                             
                             # Setup training
                             optimizer = AdamW(self._model.parameters(), lr=2e-5)
                             self._model.train()
-                            logger.info("Starting training loop...")
+                            logger.info("Starting training...")
                             
                             # Single example training
                             for batch in train_loader:
@@ -199,16 +160,16 @@ class BertClassifier(LabelStudioMLBase):
                                 )
                                 
                                 loss = outputs.loss
-                                logger.info(f"Training loss: {loss.item()}")
+                                logger.info(f"Training loss: {loss.item():.4f}")
                                 
                                 # Backward pass
                                 loss.backward()
                                 optimizer.step()
                             
-                            # Save the model after training
+                            # Save the model
                             model_path = os.path.join(self.model_dir, 'model_state.pt')
                             torch.save(self._model.state_dict(), model_path)
-                            logger.info(f"Model saved to {model_path}")
+                            logger.info(f"✓ Model saved to {model_path}")
                             
                             return {
                                 'status': 'ok',