Internal change

PiperOrigin-RevId: 433453559

official/utils/docs/build_nlp_api_docs.py

@@ -37,7 +37,7 @@ FLAGS = flags.FLAGS
 flags.DEFINE_string('output_dir', None, 'Where to write the resulting docs to.')
 flags.DEFINE_string(
     'code_url_prefix',
-    'https://github.com/tensorflow/models/blob/master/official/nlp/modeling/',
+    'https://github.com/tensorflow/models/blob/master/official/nlp/modeling',
     'The url prefix for links to code.')
 
 flags.DEFINE_bool('search_hints', True,

official/utils/docs/build_vision_api_docs.py

@@ -21,7 +21,7 @@ $> python build_vision_api_docs \
  --output_dir=/tmp/api_docs
 """
 
-import os
+import pathlib
 
 from absl import app
 from absl import flags
@@ -29,15 +29,16 @@ from absl import logging
 from tensorflow_docs.api_generator import generate_lib
 from tensorflow_docs.api_generator import public_api
 
+import tensorflow_models as tfm
 import build_api_docs_lib
-from official.vision.beta import modeling as tfvision
+
 
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string('output_dir', None, 'Where to write the resulting docs to.')
 flags.DEFINE_string(
     'code_url_prefix',
-    'https://github.com/tensorflow/models/blob/master/official/vision/beta/modeling/',
+    'https://github.com/tensorflow/models/blob/master/tensorflow_models/',
     'The url prefix for links to code.')
 
 flags.DEFINE_bool('search_hints', True,
@@ -46,23 +47,47 @@ flags.DEFINE_bool('search_hints', True,
 flags.DEFINE_string('site_path', 'tfvision/api_docs/python',
                     'Path prefix in the _toc.yaml')
 
-PROJECT_SHORT_NAME = 'tfvision'
+PROJECT_SHORT_NAME = 'tfm.vision'
 PROJECT_FULL_NAME = 'TensorFlow Official Models - Vision Modeling Library'
 
 
+def custom_filter(path, parent, children):
+  if len(path) == 1:
+    # Don't filter the contents of the top level `tfm.vision` package.
+    return children
+  else:
+    return public_api.explicit_package_contents_filter(path, parent, children)
+
+
 def gen_api_docs(code_url_prefix, site_path, output_dir, project_short_name,
                  project_full_name, search_hints):
   """Generates api docs for the tensorflow docs package."""
   build_api_docs_lib.hide_module_model_and_layer_methods()
 
+  branch = code_url_prefix.strip('/').split('/')[-2]
+  official_url_prefix = (
+      f'https://github.com/tensorflow/models/blob/{branch}/official/')
+
+  vision_base_dir = pathlib.Path(tfm.vision.__file__).parent
+
+  # The `layers` submodule (and others) are actually defined in the `official`
+  # package. Find the path to `official`.
+  official_base_dir = [
+      p for p in pathlib.Path(tfm.vision.layers.__file__).parents
+      if p.name == 'official'
+  ][0]
+
   doc_generator = generate_lib.DocGenerator(
       root_title=project_full_name,
-      py_modules=[(project_short_name, tfvision)],
-      base_dir=os.path.dirname(tfvision.__file__),
-      code_url_prefix=code_url_prefix,
+      py_modules=[(project_short_name, tfm.vision)],
+      base_dir=[
+          vision_base_dir,
+          official_base_dir,
+      ],
+      code_url_prefix=[code_url_prefix, official_url_prefix],
       search_hints=search_hints,
       site_path=site_path,
-      callbacks=[public_api.explicit_package_contents_filter],
+      callbacks=[custom_filter],
   )
 
   doc_generator.build(output_dir)