making it atrt training

model.py

@@ -142,54 +142,88 @@ class BertClassifier(LabelStudioMLBase):
         return predictions
 
     def fit(self, event_data, data=None, **kwargs):
-        """Train the model on the labeled data."""
-        logger.info("=== WEBHOOK DEBUG INFO ===")
-        logger.info(f"event_data type: {type(event_data)}")
-        logger.info(f"event_data content: {event_data}")
-        logger.info(f"data type: {type(data)}")
-        logger.info(f"data content: {data}")
-        logger.info(f"kwargs: {kwargs}")
-        logger.info("=== END WEBHOOK DEBUG INFO ===")
-        
-        logger.info(f"Received event: {event_data}")
+        """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)}")
         
         try:
             if event_data == 'ANNOTATION_CREATED':
-                # Extract text and label directly from the data
+                logger.info("Processing ANNOTATION_CREATED event")
                 annotation = data.get('annotation', {})
                 task = data.get('task', {})
                 
+                logger.info(f"Annotation data: {json.dumps(annotation, indent=2)}")
+                logger.info(f"Task data: {json.dumps(task, indent=2)}")
+                
                 if not task or not annotation:
                     logger.error("Missing task or annotation data")
                     return {'status': 'error', 'message': 'Missing task or annotation data'}
                 
-                # Get the text from task data
+                # Extract text and label
                 text = task.get('data', {}).get('text', '')
-                
-                # Get the label from annotation results
                 results = annotation.get('result', [])
+                
                 for result in results:
                     if result.get('type') == 'choices':
                         label = result.get('value', {}).get('choices', [])[0]
-                        logger.info(f"Processing annotation - Text: {text[:50]}... Label: {label}")
+                        logger.info(f"Training on - Text: {text[:50]}... Label: {label}")
                         
-                        # Here you would add your training logic
-                        # For now, let's just log it
-                        logger.info(f"Would train model on text: '{text}' with label: '{label}'")
-                        
-                        return {
-                            'status': 'ok',
-                            'message': f'Added training data: {text[:50]}... -> {label}'
-                        }
-            
-            elif event_data == 'START_TRAINING':
-                # This event indicates we should start a training cycle
-                logger.info("Received START_TRAINING event")
-                # Here you would implement the actual training logic
-                return {'status': 'ok', 'message': 'Training cycle started'}
-                
+                        try:
+                            # Create dataset for single example
+                            dataset = TextDataset(
+                                texts=[text],
+                                labels=[self.categories.index(label)],
+                                tokenizer=self.tokenizer
+                            )
+                            train_loader = DataLoader(dataset, batch_size=1)
+                            
+                            # Setup training
+                            optimizer = AdamW(self._model.parameters(), lr=2e-5)
+                            self._model.train()
+                            
+                            # Single example training
+                            for batch in train_loader:
+                                optimizer.zero_grad()
+                                
+                                # Move batch to device
+                                input_ids = batch['input_ids'].to(self.device)
+                                attention_mask = batch['attention_mask'].to(self.device)
+                                labels = batch['labels'].to(self.device)
+                                
+                                # Forward pass
+                                outputs = self._model(
+                                    input_ids=input_ids,
+                                    attention_mask=attention_mask,
+                                    labels=labels
+                                )
+                                
+                                loss = outputs.loss
+                                logger.info(f"Training loss: {loss.item()}")
+                                
+                                # Backward pass
+                                loss.backward()
+                                optimizer.step()
+                            
+                            # Save the model after training
+                            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}")
+                            
+                            return {
+                                'status': 'ok',
+                                'message': f'Successfully trained on: {text[:50]}... -> {label}',
+                                'time_taken': str(datetime.now() - start_time)
+                            }
+                            
+                        except Exception as e:
+                            logger.error(f"Training error: {str(e)}")
+                            logger.error("Full error details:", exc_info=True)
+                            return {'status': 'error', 'message': f'Training failed: {str(e)}'}
+    
         except Exception as e:
-            logger.error(f"Error during training: {str(e)}")
+            logger.error(f"Error in fit method: {str(e)}")
             logger.error("Full error details:", exc_info=True)
             return {'status': 'error', 'message': str(e)}