Merge pull request #136 from georgian-io/jsonl-support

`jsonl` format support

Author: benjaminye

llmtune/data/ingestor.py

@@ -8,12 +8,14 @@
 def get_ingestor(data_type: str):
     if data_type == "json":
         return JsonIngestor
+    elif data_type == "jsonl":
+        return JsonlIngestor
     elif data_type == "csv":
         return CsvIngestor
     elif data_type == "huggingface":
         return HuggingfaceIngestor
     else:
-        raise ValueError(f"'type' must be one of 'json', 'csv', or 'huggingface', you have {data_type}")
+        raise ValueError(f"'type' must be one of 'json', 'jsonl', 'csv', or 'huggingface', you have {data_type}")
 
 
 class Ingestor(ABC):
@@ -35,6 +37,19 @@ def to_dataset(self) -> Dataset:
         return Dataset.from_generator(self._json_generator)
 
 
+class JsonlIngestor(Ingestor):
+    def __init__(self, path: str):
+        self.path = path
+
+    def _jsonl_generator(self):
+        with open(self.path, "rb") as f:
+            for item in ijson.items(f, "", multiple_values=True):
+                yield item
+
+    def to_dataset(self) -> Dataset:
+        return Dataset.from_generator(self._jsonl_generator)
+
+
 class CsvIngestor(Ingestor):
     def __init__(self, path: str):
         self.path = path

llmtune/pydantic_models/config_model.py

@@ -13,7 +13,7 @@ class QaConfig(BaseModel):
 
 
 class DataConfig(BaseModel):
-    file_type: Literal["json", "csv", "huggingface"] = Field(None, description="File type")
+    file_type: Literal["json", "jsonl", "csv", "huggingface"] = Field(None, description="File type")
     path: Union[FilePath, HfModelPath] = Field(None, description="Path to the file or HuggingFace model")
     prompt: str = Field(None, description="Prompt for the model. Use {} brackets for column name")
     prompt_stub: str = Field(