diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx
index 609899aa95a..84b60f1e039 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx
@@ -28,16 +28,12 @@ const PiiConfiguration: React.FC<PiiConfigurationProps> = ({
       ? [...entityCategories, { category: "Custom", entities: customEntities }]
       : entityCategories;
 
-  // Create a lookup map to quickly find an entity's category
-  const entityToCategoryMap = new Map<string, string>();
-  entityCategories.forEach((category) => {
-    category.entities.forEach((entity) => {
-      entityToCategoryMap.set(entity, category.category);
-    });
-  });
-  customEntities.forEach((entity) => {
-    entityToCategoryMap.set(entity, "Custom");
-  });
+  const entityToCategoryMap = new Map<string, string>([
+    ...entityCategories.flatMap((category) =>
+      category.entities.map((entity): [string, string] => [entity, category.category]),
+    ),
+    ...customEntities.map((entity): [string, string] => [entity, "Custom"]),
+  ]);
 
   // Filter entities based on selected categories
   const filteredEntities = allEntities.filter((entity) => {