distillation: fix ModuleNotFoundError error in token counts script
Hi,

I'm currently trying out "distillation" 😅

This PR a `ModuleNotFoundError` message in the `token_counts.py` script (same error was recently fixed in 803c1cc4eacd38f1b854578d7d717b5e4a1ada47 🤗
Great, thanks @stefan-it!