diff --git a/docfx_yaml/extension.py b/docfx_yaml/extension.py
index 97af8ca8..d55219bc 100644
--- a/docfx_yaml/extension.py
+++ b/docfx_yaml/extension.py
@@ -1272,13 +1272,13 @@ def convert_module_to_package_if_needed(obj):
                         del(obj['source'])
 
             yaml_map[uid] = [yaml_data, references]
-           
+
             # Parse the name of the object.
             # Some types will need additional parsing to de-duplicate their names and contain
-            # a portion of their parent name for better disambiguation. This is done in 
+            # a portion of their parent name for better disambiguation. This is done in
             # disambiguate_toc_name
-            
-            node_name = obj.get('class').split(".")[-1] if obj.get('class') else obj['name']
+
+            node_name = uid.split(".")[-1]
 
             # Build nested TOC
             if uid.count('.') >= 1:
@@ -1286,38 +1286,25 @@ def convert_module_to_package_if_needed(obj):
                 found_node = find_node_in_toc_tree(pkg_toc_yaml, parent_level)
 
                 if found_node:
-                    found_node.pop('uid', 'No uid found')
                     found_node.setdefault(
-                      'items', 
+                      'items',
                       [{'name': 'Overview', 'uidname': parent_level, 'uid': parent_level}]
                     ).append({
                       'name': node_name,
-                      'uidname': uid, 
+                      'uidname': uid,
                       'uid': uid
                     })
                 else:
-                    # Extract the file name to use for new entries into the TOC
-                    # to make TOC entries consistent with filenames.
-                    file_name = obj['source']['path']
-                    # `inspect.getfile() cannot retrieve file path for some
-                    # types. Try to retrieve the file path from its class
-                    # parent, otherwise simply use node_name.
-                    if not file_name:
-                        file_name = app.env.docfx_class_paths.get(obj.get("class"))
-                    try:
-                        file_name = obj['source']['path'].split("/")[-1][:-3]
-                    except AttributeError:
-                        file_name = node_name
                     pkg_toc_yaml.append({
-                      'name': file_name, 
-                      'uidname': uid, 
+                      'name': node_name,
+                      'uidname': uid,
                       'uid': uid
                     })
 
             else:
                 pkg_toc_yaml.append({
-                  'name': node_name, 
-                  'uidname': uid, 
+                  'name': node_name,
+                  'uidname': uid,
                   'uid': uid
                 })
 
diff --git a/tests/test_unit.py b/tests/test_unit.py
index e7279de6..39443bb1 100644
--- a/tests/test_unit.py
+++ b/tests/test_unit.py
@@ -17,7 +17,7 @@ class TestGenerate(unittest.TestCase):
     def test_find_unique_name(self):
 
         entries = {}
-        
+
         # Disambiguate with unique entries.
         entry1 = "google.cloud.aiplatform.v1.schema.predict.instance_v1.types"
         entry2 = "google.cloud.aiplatform.v1beta2.schema.predict.instance_v1.types"