From 58b6d985c8e57d14d0b7fb1eeddcea42105b028e Mon Sep 17 00:00:00 2001
From: Faiz Surani <faiz.surani@gmail.com>
Date: Mon, 30 Sep 2024 15:44:09 -0700
Subject: [PATCH] Add pydantic support to rl.utils.io.write_jsonl

---
 rl/utils/io.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/rl/utils/io.py b/rl/utils/io.py
index 7422f09..7818d71 100644
--- a/rl/utils/io.py
+++ b/rl/utils/io.py
@@ -9,6 +9,7 @@
 import dotenv
 import requests
 import tqdm
+from pydantic import BaseModel
 
 DOTENV_LOADED = False
 
@@ -98,7 +99,12 @@ def write_jsonl(filename: str | Path, records: Iterable[Any], overwrite=False) -
         raise ValueError(f"{filename} already exists and overwrite is not set.")
     with filename.open("w") as f:
         for record in records:
-            f.write(json.dumps(record) + "\n")
+            json_record: str
+            if isinstance(record, BaseModel):
+                json_record = record.model_dump_json()
+            else:
+                json_record = json.dumps(record)
+            f.write(json_record + "\n")
 
 
 def write_jsonl_spark(filename: str | Path, df, overwrite=False) -> None: