WIP

1 files changed, 4 insertions, 4 deletions

diff --git a/make_dataset.py b/make_dataset.py
index b7695b3..b2f2cee 100755
--- a/make_dataset.py
+++ b/make_dataset.py
@@ -6,7 +6,7 @@ import os
 
 from datasets import Dataset, DatasetDict
 from glob import glob
-from PIL.Image import Image
+from PIL import Image
 
 
 def make_dataset(base="./baseline"):
@@ -14,13 +14,13 @@ def make_dataset(base="./baseline"):
 
     user_prompts = []
     images = []
-    file_ids = []
+    ids = []
     assistant_replies = []
 
     for filename in glob(f"{base}/*.jsonl"):
         with open(filename, "r") as f:
             data = json.loads(f.read())
-        id_ = os.path.basename(filename)
+        id_ = os.path.basename(filename).removesuffix(".jsonl")
         image_path = f"../d/Images/{id_.removesuffix(".jsonl")}.jpg"
 
         user_prompts.append(prompt)
@@ -29,7 +29,7 @@ def make_dataset(base="./baseline"):
             "German": data["Translation"],
         }, ensure_ascii=False, indent=0))
         ids.append(id_)
-        images.append(Image(image_path).convert("RGB"))
+        images.append(Image.open(image_path).convert("RGB"))
 
     return Dataset.from_dict({"id": ids, "image": images, "user": user_prompts, "assistant": assistant_replies})