From a8dedf63c66f10dd325f79ab6db0ce56207db40b Mon Sep 17 00:00:00 2001
From: giom <giom@google.com>
Date: Wed, 1 Jan 2025 12:37:24 +0000
Subject: [PATCH] adding the missing part

---
 .github/workflows/notebooks.yaml | 60 +++++++++++++++++++++++++++++++-
 1 file changed, 59 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/notebooks.yaml b/.github/workflows/notebooks.yaml
index 1d8fdc72c..d963c36c8 100644
--- a/.github/workflows/notebooks.yaml
+++ b/.github/workflows/notebooks.yaml
@@ -19,4 +19,62 @@ jobs:
         # Checkout PR branch for PR events, main branch for manual runs
         ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || 'main' }}
         fetch-depth: 0  # Fetch all history to avoid issues with rebasing
-    - uses: actions/setup
\ No newline at end of file
+    - uses: actions/setup
+        - uses: actions/setup-python@v4
+    - name: Install tensorflow-docs
+      run: python3 -m pip install -U git+https://github.com/tensorflow/docs
+    - name: Fetch main branch
+      run: git fetch -u origin main:main
+    - name: Check notebook formatting
+      run: |
+        if [ "${{ github.event_name }}" == "pull_request" ]; then
+          # Only check notebooks modified in this pull request
+          readarray -t changed_notebooks < <(git diff --name-only main | grep '\.ipynb$' || true)
+        else
+          # Manual run, check everything
+          readarray -t changed_notebooks < <(find -name '*.ipynb')
+        fi
+        if [[ ${#changed_notebooks[@]} == 0 ]]; then
+          echo "No notebooks modified in this pull request."
+          exit 0
+        else
+          echo "Check formatting with nbfmt:"
+          python3 -m tensorflow_docs.tools.nbfmt --test "${changed_notebooks[@]}"
+        fi
+  nblint:
+    name: Notebook lint
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+    - uses: actions/setup-python@v4
+    - name: Install tensorflow-docs
+      run: python3 -m pip install -U git+https://github.com/tensorflow/docs
+    - name: Fetch main branch
+      run: git fetch -u origin main:main
+
+    # Lint for all notebooks
+    - name: Lint notebooks
+      run: |
+        if [ "${{ github.event_name }}" == "pull_request" ]; then
+          # Only check notebooks modified in this pull request
+          readarray -t changed_notebooks < <(git diff --name-only main |grep '\.ipynb$' || true)
+        else
+          # Manual run, check everything
+          readarray -t changed_notebooks < <(find . -name '*.ipynb')
+        fi
+        if [[ ${#changed_notebooks[@]} == 0 ]]; then
+          echo "No website notebooks modified in this pull request."
+          exit 0
+        else
+          echo "WARNING: If the button_colab check fails for you, make sure you have <table class=\"tfo-notebook-buttons\"...>"
+          echo "Lint check with nblint:"
+          python3 -m tensorflow_docs.tools.nblint \
+            --styles=google,tensorflow \
+            --arg=repo:google-gemini/cookbook \
+            --arg=branch:main \
+            --exclude_lint=tensorflow::button_download \
+            --exclude_lint=tensorflow::button_website \
+            --arg=base_url:https://ai.google.dev/ \
+            --exclude_lint=tensorflow::button_github \
+            "${changed_notebooks[@]}"
+        fi
\ No newline at end of file