From 3e0fe3cf5c1059c04535de8f04f4efed7251adbe Mon Sep 17 00:00:00 2001
From: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Date: Mon, 10 Aug 2020 11:12:29 -0400
Subject: [PATCH] Colab button (#6389)

* Add colab button

* Add colab link for tutorials
---
 docs/source/_static/css/huggingface.css | 31 +++++++++++++++++++++++++
 docs/source/_static/js/custom.js        | 31 +++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/docs/source/_static/css/huggingface.css b/docs/source/_static/css/huggingface.css
index c6ad7b7d5c..081e99f165 100644
--- a/docs/source/_static/css/huggingface.css
+++ b/docs/source/_static/css/huggingface.css
@@ -1,5 +1,36 @@
 /* Our DOM objects */
 
+/* Colab dropdown */
+
+.colab-dropdown {
+    position: relative;
+    display: inline-block;
+}
+  
+.colab-dropdown-content {
+    display: none;
+    position: absolute;
+    background-color: #f9f9f9;
+    min-width: 117px;
+    box-shadow: 0px 8px 16px 0px rgba(0,0,0,0.2);
+    z-index: 1;
+}
+  
+.colab-dropdown-content button {
+    color: #6670FF;
+    background-color: #f9f9f9;
+    font-size: 12px;
+    border: none;
+    min-width: 117px;
+    padding: 5px 5px;
+    text-decoration: none;
+    display: block;
+}
+  
+.colab-dropdown-content button:hover {background-color: #eee;}
+  
+.colab-dropdown:hover .colab-dropdown-content {display: block;}
+
 /* Version control */
 
 .version-button {
diff --git a/docs/source/_static/js/custom.js b/docs/source/_static/js/custom.js
index 1ce992ddfe..0aeec4f6ec 100644
--- a/docs/source/_static/js/custom.js
+++ b/docs/source/_static/js/custom.js
@@ -21,6 +21,17 @@ const versionMapping = {
     "v1.1.0": "v1.1.0",
     "v1.0.0": "v1.0.0"
 }
+// The page that have a notebook and therefore should have the open in colab badge.
+const hasNotebook = [
+    "benchmarks",
+    "multilingual",
+    "perplexity",
+    "preprocessing",
+    "quicktour",
+    "task_summary",
+    "tokenizer_summary",
+    "training"
+];
 
 function addIcon() {
     const huggingFaceLogo = "https://huggingface.co/landing/assets/transformers-docs/huggingface_logo.svg";
@@ -82,6 +93,25 @@ function addGithubButton() {
     document.querySelector(".wy-side-nav-search .icon-home").insertAdjacentHTML('afterend', div);
 }
 
+function addColabLink() {
+    const parts = location.toString().split('/');
+    const pageName = parts[parts.length - 1].split(".")[0];
+
+    if (hasNotebook.includes(pageName)) {
+        const linksColab = `
+        <div class="colab-dropdown">
+            <img alt="Open In Colab" src="https://colab.research.google.com/assets/colab-badge.svg">
+            <div class="colab-dropdown-content">
+                <button onclick=" window.open('https://colab.research.google.com/github/sgugger/test_nbs/blob/master/${pageName}.ipynb')">Mixed</button>
+                <button onclick=" window.open('https://colab.research.google.com/github/sgugger/test_nbs/blob/master/pytorch/${pageName}.ipynb')">PyTorch</button>
+                <button onclick=" window.open('https://colab.research.google.com/github/sgugger/test_nbs/blob/master/tensorflow/${pageName}.ipynb')">TensorFlow</button>
+            </div>
+        </div>`
+        const leftMenu = document.querySelector(".wy-breadcrumbs-aside")
+        leftMenu.innerHTML = linksColab + '\n' + leftMenu.innerHTML
+    }
+}
+
 function addVersionControl() {
     // To grab the version currently in view, we parse the url
     const parts = location.toString().split('/');
@@ -255,6 +285,7 @@ function onLoad() {
     addGithubButton();
     parseGithubButtons();
     addHfMenu();
+    addColabLink();
     platformToggle();
 }