Define defaults for ChunkTask

Aleph-Alpha · Feb 29, 2024 · 4be3230 · 4be3230
1 parent d4db4b7
commit 4be3230
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/intelligence_layer/core/chunk.py b/src/intelligence_layer/core/chunk.py
@@ -3,7 +3,7 @@
 from pydantic import BaseModel
 from semantic_text_splitter import HuggingFaceTextSplitter
 
-from intelligence_layer.core.model import ControlModel
+from intelligence_layer.core.model import ControlModel, LuminousControlModel
 from intelligence_layer.core.task import Task
 from intelligence_layer.core.tracer.tracer import TaskSpan
 
@@ -50,8 +50,11 @@ class ChunkTask(Task[ChunkInput, ChunkOutput]):
         max_tokens_per_chunk: The maximum number of tokens to fit into one chunk.
     """
 
-    def __init__(self, model: ControlModel, max_tokens_per_chunk: int):
+    def __init__(
+        self, model: ControlModel | None = None, max_tokens_per_chunk: int = 512
+    ):
         super().__init__()
+        model = model or LuminousControlModel()
         self._splitter = HuggingFaceTextSplitter(model.get_tokenizer())
         self._max_tokens_per_chunk = max_tokens_per_chunk