diff --git a/docs/css/api-docs.css b/docs/css/api-docs.css
index cc5f37bbdb42e..b2d1d7f869790 100644
--- a/docs/css/api-docs.css
+++ b/docs/css/api-docs.css
@@ -8,6 +8,10 @@
   background-color: #257080;
 }
 
+.alphaComponent {
+  background-color: #bb0000;
+}
+
 .badge {
   font-family: Arial, san-serif;
   float: right;
diff --git a/docs/js/api-docs.js b/docs/js/api-docs.js
index 2c3ca434ec46c..1414b6d0b81a1 100644
--- a/docs/js/api-docs.js
+++ b/docs/js/api-docs.js
@@ -2,16 +2,14 @@
 
 $(document).ready(function() {
   var annotations = $("dt:contains('Annotations')").next("dd").children("span.name");
-  addBadges(annotations, "AlphaComponent", ":: AlphaComponent ::", "<span class='experimental badge'>ALPHA COMPONENT</span>");
+  addBadges(annotations, "AlphaComponent", ":: AlphaComponent ::", "<span class='alphaComponent badge'>Alpha Component</span>");
   addBadges(annotations, "DeveloperApi", ":: DeveloperApi ::", "<span class='developer badge'>Developer API</span>");
   addBadges(annotations, "Experimental", ":: Experimental ::", "<span class='experimental badge'>Experimental</span>");
 });
 
 function addBadges(allAnnotations, name, tag, html) {
-  var fullName = "org.apache.spark.annotations." + name;
-  var annotations = allAnnotations.children("a[name='" + fullName + "']");
-  var tags = $("p.comment:contains(" + tag + ")").add(
-    $("div.comment p:contains(" + tag + ")"));
+  var annotations = allAnnotations.filter(":contains('" + name + "')")
+  var tags = $(".cmt:contains(" + tag + ")")
 
   // Remove identifier tags from comments
   tags.each(function(index) {
@@ -21,7 +19,8 @@ function addBadges(allAnnotations, name, tag, html) {
   });
 
   // Add badges to all containers
-  tags.prevAll("h4.signature").prepend(html);
-  annotations.closest("div.fullcomment").prevAll("h4.signature").prepend(html);
-  annotations.closest("div.fullcommenttop").prepend(html);
+  tags.prevAll("h4.signature")
+    .add(annotations.closest("div.fullcommenttop"))
+    .add(annotations.closest("div.fullcomment").prevAll("h4.signature"))
+    .prepend(html);
 }