Merge pull request #119 from sbintuitions/smart_open

Use `smart_open` for jsonl loading
sbintuitions · Jan 9, 2025 · 77d6a07 · 77d6a07
2 parents f9b260a + 8661c70
commit 77d6a07
Show file tree

Hide file tree

Showing 6 changed files with 106 additions and 1 deletion.
diff --git a/flexeval/core/chat_dataset/template_based.py b/flexeval/core/chat_dataset/template_based.py
@@ -6,6 +6,7 @@
 
 import datasets
 from jinja2 import Template
+from smart_open import open
 
 from flexeval.core.utils.jinja2_utils import JINJA2_ENV
 

diff --git a/flexeval/core/generation_dataset/template_based.py b/flexeval/core/generation_dataset/template_based.py
@@ -6,6 +6,7 @@
 
 import datasets
 from jinja2 import Template
+from smart_open import open
 
 from flexeval.core.utils.jinja2_utils import JINJA2_ENV
 

diff --git a/flexeval/core/reward_bench_dataset/template_based.py b/flexeval/core/reward_bench_dataset/template_based.py
@@ -5,6 +5,7 @@
 
 import datasets
 from jinja2 import Template
+from smart_open import open
 
 from flexeval.core.utils.jinja2_utils import JINJA2_ENV
 

diff --git a/flexeval/core/text_dataset/jsonl.py b/flexeval/core/text_dataset/jsonl.py
@@ -3,6 +3,8 @@
 import json
 from os import PathLike
 
+from smart_open import open
+
 from .base import TextDataset, TextInstance
 
 

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -34,6 +34,7 @@ loguru = "^0.7.2"
 wandb = {version = "^0.17.2", optional = true}
 pyarrow = "16.1.0"  # set the version because we get "Unable to find installation candidates" with 17.0.0
 scipy = "1.13.0"
+smart-open = "^7.1.0"
 
 [tool.poetry.extras]
 vllm = ["vllm"]