Epochs added

docker-compose.yml

@@ -20,6 +20,7 @@ services:
       - LORA_DROPOUT=0.1
       - LORA_TARGET_MODULES=q,v
       # Training settings
+      - EPOCHS=3
       - LEARNING_RATE=1e-4
       - BATCH_SIZE=1
       - MAX_STEPS=100

model.py

@@ -207,8 +207,8 @@ class T5Model(LabelStudioMLBase):
             
             # Configure LoRA
             lora_config = LoraConfig(
-                r=int(os.getenv('LORA_R', '8')),
-                lora_alpha=int(os.getenv('LORA_ALPHA', '32')),
+                r=int(os.getenv('LORA_R', '16')),
+                lora_alpha=int(os.getenv('LORA_ALPHA', '16')),
                 target_modules=os.getenv('LORA_TARGET_MODULES', 'q,v').split(','),
                 lora_dropout=float(os.getenv('LORA_DROPOUT', '0.1')),
                 bias="none",
@@ -227,16 +227,20 @@ class T5Model(LabelStudioMLBase):
             logger.info("Starting training loop...")
             optimizer = torch.optim.AdamW(model.parameters(), lr=float(os.getenv('LEARNING_RATE', '1e-4')))
             
-            # Set model to training mode
-            model.train()
-            optimizer.zero_grad()
+            num_epochs = int(os.getenv('NUM_EPOCHS', '3'))
             
-            outputs = model(**inputs, labels=labels["input_ids"])
-            loss = outputs.loss
-            loss.backward()
-            optimizer.step()
-            
-            logger.info(f"Training step completed. Loss: {loss.item():.4f}")
+            for epoch in range(num_epochs):
+                logger.info(f"Starting epoch {epoch+1}/{num_epochs}")
+                
+                model.train()
+                optimizer.zero_grad()
+                
+                outputs = model(**inputs, labels=labels["input_ids"])
+                loss = outputs.loss
+                loss.backward()
+                optimizer.step()
+                
+                logger.info(f"Epoch {epoch+1}/{num_epochs} completed. Loss: {loss.item():.4f}")
             
             # Switch back to eval mode
             model.eval()