Update app.py

app.py

@@ -178,36 +178,37 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
         decoded_text = [tokenizer.decode(tokens, skip_special_tokens=True).replace(user_msg, "") for tokens in generated_response]
 
         rewards = list()
-        toxicities = list()
+        
+        if safety == "On":
+            toxicities = list()
 
         for text in decoded_text:
-          reward_tokens = rewardTokenizer(user_msg, text,
-                        truncation=True,
-                        max_length=512,
-                        return_token_type_ids=False,
-                        return_tensors="pt",
-                        return_attention_mask=True)
-
-          reward_tokens.to(rewardModel.device)
-
-          reward = rewardModel(**reward_tokens)[0].item()
-
-          toxicity_tokens = toxiciyTokenizer(user_msg + " " + text,
+            reward_tokens = rewardTokenizer(user_msg, text,
                         truncation=True,
                         max_length=512,
                         return_token_type_ids=False,
                         return_tensors="pt",
                         return_attention_mask=True)
-
-          toxicity_tokens.to(toxicityModel.device)
-
-          toxicity = toxicityModel(**toxicity_tokens)[0].item()
-
-          rewards.append(reward)
-          toxicities.append(toxicity)
+            
+            reward_tokens.to(rewardModel.device)
+            
+            reward = rewardModel(**reward_tokens)[0].item()
+            rewards.append(reward)
+
+            if safety == "On":
+                toxicity_tokens = toxiciyTokenizer(user_msg + " " + text,
+                            truncation=True,
+                            max_length=512,
+                            return_token_type_ids=False,
+                            return_tensors="pt",
+                            return_attention_mask=True)
+                
+                toxicity_tokens.to(toxicityModel.device)
+                
+                toxicity = toxicityModel(**toxicity_tokens)[0].item()
+                toxicities.append(toxicity)
+                toxicity_threshold = 5
         
-        toxicity_threshold = 5
-
         ordered_generations = sorted(zip(decoded_text, rewards, toxicities), key=lambda x: x[1], reverse=True)
 
         if safety == "On":