Adding training initiation and logging

model.py

@@ -7,6 +7,8 @@ import json
 import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from label_studio_ml.model import LabelStudioMLBase, ModelResponse
+from peft import get_peft_model, LoraConfig
+import time
 
 logger = logging.getLogger(__name__)
 
@@ -180,26 +182,61 @@ Category:"""
         return predictions
 
     def fit(self, event, data, **kwargs):
-        """Handle annotation events from Label Studio
-        
-        Args:
-            event (str): Event type ('ANNOTATION_CREATED', 'ANNOTATION_UPDATED', 'START_TRAINING')
-            data (dict): Event payload with annotation details
-        """
+        """Handle annotation events from Label Studio"""
+        start_time = time.time()
+        logger.info("Starting training session...")
         
         valid_events = {'ANNOTATION_CREATED', 'ANNOTATION_UPDATED', 'START_TRAINING'}
         if event not in valid_events:
             logger.warning(f"Skip training: event {event} is not supported")
             return
         
-        # Extract text and label
-        text = data['task']['data']['text']
-        label = data['annotation']['result'][0]['value']['choices'][0]
-        
-        training_data = {
-            'text': text,
-            'label': label
-        }
-        
-        logger.info(f"Extracted training data: {json.dumps(training_data, indent=2)}")
+        try:
+            # Extract text and label
+            text = data['task']['data']['text']
+            label = data['annotation']['result'][0]['value']['choices'][0]
+            
+            # Configure LoRA
+            lora_config = LoraConfig(
+                r=int(os.getenv('LORA_R', '8')),
+                lora_alpha=int(os.getenv('LORA_ALPHA', '32')),
+                target_modules=os.getenv('LORA_TARGET_MODULES', 'q,v').split(','),
+                lora_dropout=float(os.getenv('LORA_DROPOUT', '0.1')),
+                bias="none",
+                task_type="SEQ_2_SEQ_LM"
+            )
+            
+            logger.info("Preparing model for training...")
+            model = get_peft_model(self.model, lora_config)
+            model.print_trainable_parameters()
+            
+            # Training loop
+            logger.info("Starting training loop...")
+            optimizer = torch.optim.AdamW(model.parameters(), lr=float(os.getenv('LEARNING_RATE', '1e-4')))
+            
+            # Single training step for this annotation
+            model.train()
+            optimizer.zero_grad()
+            
+            inputs = self.tokenizer(text, return_tensors="pt", max_length=self.max_length, truncation=True).to(self.device)
+            labels = self.tokenizer(label, return_tensors="pt", max_length=self.generation_max_length, truncation=True).to(self.device)
+            
+            outputs = model(**inputs, labels=labels["input_ids"])
+            loss = outputs.loss
+            loss.backward()
+            optimizer.step()
+            
+            logger.info(f"Training step completed. Loss: {loss.item():.4f}")
+            
+            # Switch back to eval mode
+            model.eval()
+            
+            training_time = time.time() - start_time
+            logger.info(f"Training session completed successfully in {training_time:.2f} seconds with tag: '{text}' and label: '{label}'")
+            
+        except Exception as e:
+            training_time = time.time() - start_time
+            logger.error(f"Training failed after {training_time:.2f} seconds")
+            logger.error(f"Error during training: {str(e)}")
+            raise