Update app.py

app.py

@@ -6,10 +6,17 @@ from opacus import PrivacyEngine
 from torch.utils.data import Dataset
 from torch.optim import AdamW
 
+# Disable torch.compile to avoid meta device issues
+torch._dynamo.config.suppress_errors = True
+torch.set_default_dtype(torch.float32)
+
+# Set device explicitly
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # Load ParsBERT model and tokenizer
 model_name = "HooshvareLab/bert-base-parsbert-uncased"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = BertLMHeadModel.from_pretrained(model_name, is_decoder=True)
+model = BertLMHeadModel.from_pretrained(model_name, is_decoder=True).to(device)
 
 # Differential Privacy parameters
 epsilon = 1.0  # Privacy budget
@@ -29,7 +36,7 @@ class ChatDataset(Dataset):
         return len(self.data)
 
     def __getitem__(self, idx):
-        item = {key: val[idx] for key, val in self.encodings.items()}
+        item = {key: val[idx].to(device) for key, val in self.encodings.items()}
         item["labels"] = item["input_ids"].clone()
         return item
 
@@ -80,23 +87,26 @@ def train_model():
     model.train()
 
     # Add differential privacy
-    privacy_engine = PrivacyEngine(secure_mode=False)  # False for experimentation
-    private_model, private_optimizer, train_dataloader = privacy_engine.make_private(
-        module=model,
-        optimizer=optimizer,
-        data_loader=trainer.get_train_dataloader(),
-        noise_multiplier=1.1,
-        max_grad_norm=1.0,
-    )
+    try:
+        privacy_engine = PrivacyEngine(secure_mode=False)  # False for experimentation
+        private_model, private_optimizer, train_dataloader = privacy_engine.make_private(
+            module=model,
+            optimizer=optimizer,
+            data_loader=trainer.get_train_dataloader(),
+            noise_multiplier=1.1,
+            max_grad_norm=1.0,
+        )
 
-    trainer.optimizer = private_optimizer
-    trainer.train_dataloader = train_dataloader
-    trainer.model = private_model  # Update trainer to use private model
+        trainer.optimizer = private_optimizer
+        trainer.train_dataloader = train_dataloader
+        trainer.model = private_model  # Update trainer to use private model
 
-    trainer.train()
-    model.save_pretrained("./fine_tuned_model")
-    tokenizer.save_pretrained("./fine_tuned_model")
-    print("Model training completed and saved to ./fine_tuned_model")
+        trainer.train()
+        model.save_pretrained("./fine_tuned_model")
+        tokenizer.save_pretrained("./fine_tuned_model")
+        print("Model training completed and saved to ./fine_tuned_model")
+    except Exception as e:
+        print(f"Error during training with differential privacy: {e}")
 
 def add_noise(tensor, sensitivity, epsilon, delta):
     """Add Laplace noise for differential privacy."""
@@ -117,7 +127,7 @@ def chat(message, history):
     model.eval()
     
     # Tokenize input
-    inputs = tokenizer(message, return_tensors="pt", padding=True, truncation=True, max_length=128).to(model.device)
+    inputs = tokenizer(message, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
     
     # Generate response with model using beam search
     with torch.no_grad():