Update setup.py

Running `main.py` yields the following:

```
ImportError: This modeling file requires the following packages that were not found in your environment: einops. Run `pip install einops`
```

This commit should fix this.

setup.py

@@ -26,6 +26,7 @@ setuptools.setup(
     python_requires=">=3.9",
     install_requires=[
         "datasets>=2.0.0",
+        "einops",
         "jsonlines",
         "numexpr",
         "openai>=0.6.4",