Merge pull request #119 from sbintuitions/smart_open

ryokan0123 · web-flow · commit 77d6a07b9086 · 2025-01-09T15:52:43.000+09:00
Use `smart_open` for jsonl loading
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"]
