moving LoRA settings under .yml

docker-compose.yml

@@ -15,13 +15,13 @@ services:
       - GENERATION_MAX_LENGTH=128
       - NUM_RETURN_SEQUENCES=1
       # LoRA settings
-      - LORA_R=8
-      - LORA_ALPHA=32
+      - LORA_R=4
+      - LORA_ALPHA=8
       - LORA_DROPOUT=0.1
       - LORA_TARGET_MODULES=q,v
       # Training settings
-      - EPOCHS=3
-      - LEARNING_RATE=1e-4
+      - EPOCHS=6
+      - LEARNING_RATE=1e-5
       - BATCH_SIZE=1
       - MAX_STEPS=100
       - SAVE_STEPS=50

model.py

@@ -207,10 +207,10 @@ class T5Model(LabelStudioMLBase):
             
             # Configure LoRA
             lora_config = LoraConfig(
-                r=int(os.getenv('LORA_R', '4')),
-                lora_alpha=int(os.getenv('LORA_ALPHA', '8')),
-                target_modules=os.getenv('LORA_TARGET_MODULES', 'q,v').split(','),
-                lora_dropout=float(os.getenv('LORA_DROPOUT', '0.1')),
+                r=int(os.getenv('LORA_R')),
+                lora_alpha=int(os.getenv('LORA_ALPHA')),
+                target_modules=os.getenv('LORA_TARGET_MODULES').split(','),
+                lora_dropout=float(os.getenv('LORA_DROPOUT')),
                 bias="none",
                 task_type="SEQ_2_SEQ_LM"
             )
@@ -225,9 +225,9 @@ class T5Model(LabelStudioMLBase):
             
             # Training loop
             logger.info("Starting training loop...")
-            optimizer = torch.optim.AdamW(model.parameters(), lr=float(os.getenv('LEARNING_RATE', '1e-5')))
+            optimizer = torch.optim.AdamW(model.parameters(), lr=float(os.getenv('LEARNING_RATE')))
             
-            num_epochs = int(os.getenv('NUM_EPOCHS', '6'))
+            num_epochs = int(os.getenv('NUM_EPOCHS'))
             
 
             # Add LoRA settings logging here