Update app.py

app.py

@@ -62,7 +62,7 @@ def train_model():
         per_device_train_batch_size=4,  # Reduced to avoid memory issues
         save_steps=10_000,
         save_total_limit=2,
-        use_cpu=True if not torch.cuda.is_available() else False,  # Replace no_cuda with use_cpu
+        use_cpu=True if not torch.cuda.is_available() else False,  # Use CPU if no GPU
         logging_steps=100,  # Log training progress
     )
 
@@ -76,8 +76,11 @@ def train_model():
         optimizers=(optimizer, None),  # Pass optimizer explicitly
     )
 
+    # Set model to training mode for differential privacy
+    model.train()
+
     # Add differential privacy
-    privacy_engine = PrivacyEngine(secure_mode=False)  # False for experimentation, set to True for production
+    privacy_engine = PrivacyEngine(secure_mode=False)  # False for experimentation
     private_model, private_optimizer, train_dataloader = privacy_engine.make_private(
         module=model,
         optimizer=optimizer,
@@ -110,6 +113,9 @@ def update_model(user_input, response):
     return f"Learning from: {user_input} -> {response}"
 
 def chat(message, history):
+    # Set model to evaluation mode for inference
+    model.eval()
+    
     # Tokenize input
     inputs = tokenizer(message, return_tensors="pt", padding=True, truncation=True, max_length=128).to(model.device)