无法从aitextgen导入TextDataset类，遇DeepSpeedPlugin导入错误求助

问题重现 #

尝试用自有职位描述数据集微调GPT-Neo模型，导入aitextgen.TokenDataset时触发以下错误：

!pip install aitextgen
import aitextgen.TokenDataset 

完整报错：

ImportError                               Traceback (most recent call last)
<ipython-input-3-24eff681443c> in <cell line: 1>()
----> 1 import aitextgen.TokenDataset

1 frames
/usr/local/lib/python3.10/dist-packages/aitextgen/aitextgen.py in <module>
     12 import torch
     13 from pkg_resources import resource_filename
---&gt; 14 from pytorch_lightning.plugins import DeepSpeedPlugin
     15 from tqdm.auto import trange
     16 from transformers import (

ImportError: cannot import name 'DeepSpeedPlugin' from 'pytorch_lightning.plugins' (/usr/local/lib/python3.10/dist-packages/pytorch_lightning/plugins/__init__.py)

问题原因 #

PyTorch Lightning 2.0+版本将DeepSpeedPlugin从pytorch_lightning.plugins迁移到了pytorch_lightning.strategies模块，而旧版本的aitextgen仍在使用旧路径导入，导致兼容性错误。

解决方案 #

方案1：安装兼容版本的PyTorch Lightning #

直接安装与旧版aitextgen兼容的PyTorch Lightning 1.x版本：

pip uninstall -y pytorch-lightning
pip install pytorch-lightning==1.9.5

安装完成后重新运行导入代码即可。

方案2：修改aitextgen源码适配新版PyTorch Lightning #

如果你不想降级PyTorch Lightning，可以手动修改aitextgen的核心文件：

1. 找到aitextgen安装目录下的aitextgen.py文件（路径类似/usr/local/lib/python3.10/dist-packages/aitextgen/aitextgen.py）
2. 将第14行的导入语句：

   from pytorch_lightning.plugins import DeepSpeedPlugin
   

   替换为：

   from pytorch_lightning.strategies import DeepSpeedStrategy
   

3. 同时检查代码中所有使用DeepSpeedPlugin的地方，替换为DeepSpeedStrategy

方案3：升级到兼容新版PyTorch Lightning的aitextgen #

检查是否有aitextgen的最新版本修复了该兼容性问题，执行升级命令：

pip install --upgrade aitextgen

升级完成后重新尝试导入。

内容的提问来源于stack exchange，提问作者nkdtech