implemented model saving

model.py

@@ -3,6 +3,7 @@ import logging
 from typing import List, Dict, Optional
 from pathlib import Path
 import json
+from datetime import datetime
 
 import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
@@ -254,6 +255,21 @@ class T5Model(LabelStudioMLBase):
                 
                 logger.info(f"Epoch {epoch+1}/{num_epochs} completed. Loss: {loss.item():.4f}")
             
+            # Save the model
+            try:
+                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                model_dir = Path(os.getenv('MODEL_DIR', '/data/models'))
+                model_dir.mkdir(parents=True, exist_ok=True)
+                
+                save_path = model_dir / f"model_{timestamp}"
+                logger.info(f"Saving model to {save_path}")
+                model.save_pretrained(save_path)
+                logger.info(f"Model successfully saved to {save_path}")
+                
+            except Exception as e:
+                logger.error(f"Failed to save model: {str(e)}")
+                raise
+            
             # Switch back to eval mode
             model.eval()