Update app.py

app.py

@@ -4,6 +4,7 @@ import numpy as np
 import gradio as gr
 from opacus import PrivacyEngine
 from torch.utils.data import Dataset
+from torch.optim import AdamW
 
 # Load ParsBERT model and tokenizer
 model_name = "HooshvareLab/bert-base-parsbert-uncased"
@@ -61,27 +62,31 @@ def train_model():
         per_device_train_batch_size=4,  # Reduced to avoid memory issues
         save_steps=10_000,
         save_total_limit=2,
-        no_cuda=True if not torch.cuda.is_available() else None,  # CPU compatibility
+        use_cpu=True if not torch.cuda.is_available() else False,  # Replace no_cuda with use_cpu
         logging_steps=100,  # Log training progress
     )
 
+    # Initialize optimizer explicitly
+    optimizer = AdamW(model.parameters(), lr=5e-5)
+
     trainer = Trainer(
         model=model,
         args=training_args,
         train_dataset=train_dataset,
+        optimizers=(optimizer, None),  # Pass optimizer explicitly
     )
 
     # Add differential privacy
-    privacy_engine = PrivacyEngine()
-    private_model, optimizer, train_dataloader = privacy_engine.make_private(
+    privacy_engine = PrivacyEngine(secure_mode=False)  # False for experimentation, set to True for production
+    private_model, private_optimizer, train_dataloader = privacy_engine.make_private(
         module=model,
-        optimizer=trainer.optimizer,
+        optimizer=optimizer,
         data_loader=trainer.get_train_dataloader(),
         noise_multiplier=1.1,
         max_grad_norm=1.0,
     )
 
-    trainer.optimizer = optimizer
+    trainer.optimizer = private_optimizer
     trainer.train_dataloader = train_dataloader
     trainer.model = private_model  # Update trainer to use private model