diff --git a/src/lean_dojo/data_extraction/ExtractData.lean b/src/lean_dojo/data_extraction/ExtractData.lean
index 0c61c5e4..0e1a5127 100644
--- a/src/lean_dojo/data_extraction/ExtractData.lean
+++ b/src/lean_dojo/data_extraction/ExtractData.lean
@@ -253,6 +253,8 @@ def findLean (mod : Name) : IO FilePath := do
     return FilePath.mk (modStr.replace "«lake-packages»" "lake-packages" |>.replace "." "/") |>.withExtension "lean"
   if modStr.startsWith "«.lake»." then
     return FilePath.mk (modStr.replace "«.lake»" ".lake" |>.replace "." "/") |>.withExtension "lean"
+  if modStr == "Lake" then
+    return packagesDir / "lean4/src/lean/lake/Lake.lean"
   let olean ← findOLean mod
   -- Remove a "build/lib/lean/" substring from the path.
   let lean := olean.toString.replace ".lake/build/lib/lean/" ""
@@ -411,7 +413,10 @@ def getImports (header: TSyntax `Lean.Parser.Module.header) : IO String := do
   let mut s := ""
 
   for dep in headerToImports header do
-    let oleanPath ← findOLean dep.module
+    -- let oleanPath ← findOLean dep.module
+    let leanPath ← Path.findLean dep.module
+    s := s ++ "\n" ++ leanPath.toString
+    /-
     if oleanPath.isRelative then
       let leanPath := Path.toSrcDir! oleanPath "lean"
       assert! ← leanPath.pathExists
@@ -429,6 +434,7 @@ def getImports (header: TSyntax `Lean.Parser.Module.header) : IO String := do
       p := p.replace "/lean4/src/lean/Lake" "/lean4/src/lean/lake/Lake"
       assert! ← FilePath.mk p |>.pathExists
       s := s ++ "\n" ++ p
+  -/
 
   return s.trim
 
diff --git a/src/lean_dojo/data_extraction/ast.py b/src/lean_dojo/data_extraction/ast.py
index 95d962b7..39e1c0e0 100644
--- a/src/lean_dojo/data_extraction/ast.py
+++ b/src/lean_dojo/data_extraction/ast.py
@@ -1401,11 +1401,10 @@ def from_data(
         start, end = None, None
         children = _parse_children(node_data, lean_file)
 
-        assert isinstance(children[0], AtomNode) and children[0].val == "import"
-        if isinstance(children[2], IdentNode):
-            module = children[2].val
-        else:
-            module = None
+        module = None
+        for child in children:
+            if isinstance(child, IdentNode):
+                module = child.val
 
         return cls(lean_file, start, end, children, module)
 
diff --git a/src/lean_dojo/data_extraction/traced_data.py b/src/lean_dojo/data_extraction/traced_data.py
index f8e47b26..837694a0 100644
--- a/src/lean_dojo/data_extraction/traced_data.py
+++ b/src/lean_dojo/data_extraction/traced_data.py
@@ -692,7 +692,10 @@ def _callback(node: Node, _):
                         if import_line.endswith(
                             suffix + ".lean"
                         ) or import_line.endswith(suffix + "/default.lean"):
-                            object.__setattr__(node, "path", Path(import_line))
+                            path = Path(import_line)
+                            if path.is_absolute():
+                                path = path.relative_to(lean_file.root_dir)
+                            object.__setattr__(node, "path", path)
 
         ast.traverse_preorder(_callback, node_cls=None)