Minor fixes on docstrings and imports