model-compose 0.3.9__tar.gz → 0.3.10__tar.gz

{model_compose-0.3.9 → model_compose-0.3.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: model-compose
-Version: 0.3.9
+Version: 0.3.10
 Summary: model-compose: Declarative AI Model and Workflow Orchestrator
 Author-email: Hanyeol Cho <hanyeol.cho@gmail.com>
 License-Expression: MIT

{model_compose-0.3.9 → model_compose-0.3.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "model-compose"
-version = "0.3.9"
+version = "0.3.10"
 description = "model-compose: Declarative AI Model and Workflow Orchestrator"
 authors = [
     { name = "Hanyeol Cho", email = "hanyeol.cho@gmail.com" }

{model_compose-0.3.9 → model_compose-0.3.10}/src/mindor/core/component/context.py

@@ -1,5 +1,5 @@
 from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, Any
-from mindor.core.utils.renderers import VariableRenderer
+from mindor.core.utils.renderers import VariableRenderer, ImageValueRenderer
 from mindor.core.gateway import find_gateway_by_port
 
 class ComponentActionContext:
@@ -13,8 +13,14 @@ class ComponentActionContext:
     def register_source(self, key: str, source: Any) -> None:
         self.sources[key] = source
 
-    async def render_variable(self, data: Dict[str, Any], ignore_files: bool = False) -> Any:
-        return await self.renderer.render(data, ignore_files)
+    def has_reference(self, key: str, value: Any) -> bool:
+        return self.renderer.has_reference(key, value)
+
+    async def render_variable(self, value: Any, ignore_files: bool = False) -> Any:
+        return await self.renderer.render(value, ignore_files)
+
+    async def render_image(self, value: Any) -> Any:
+        return await ImageValueRenderer().render(await self.render_variable(value))
 
     async def _resolve_source(self, key: str, index: Optional[int]) -> Any:
         if key in self.sources:

{model_compose-0.3.9 → model_compose-0.3.10}/src/mindor/core/component/services/model/base.py

@@ -4,7 +4,7 @@ from mindor.dsl.schema.component import ModelComponentConfig, ModelTaskType, Dev
 from mindor.dsl.schema.action import ModelActionConfig
 from mindor.core.services import AsyncService
 from ...context import ComponentActionContext
-from transformers import PreTrainedModel, PreTrainedTokenizer
+from transformers import PreTrainedModel, PreTrainedTokenizer, ProcessorMixin
 import torch, asyncio
 
 class ModelTaskService(AsyncService):
@@ -41,7 +41,7 @@ class ModelTaskService(AsyncService):
 
     def _get_common_model_params(self) -> Dict[str, Any]:
         params: Dict[str, Any] = {}
-    
+
         if self.config.device_mode != DeviceMode.SINGLE:
             params["device_map"] = self.config.device_mode.value
     
@@ -61,7 +61,10 @@ class ModelTaskService(AsyncService):
             params["local_files_only"] = True
 
         return params
-    
+
+    def _get_model_class(self) -> Type[PreTrainedModel]:
+        raise NotImplementedError("Model class loader not implemented.")
+
     def _load_pretrained_tokenizer(self, extra_params: Optional[Dict[str, Any]] = None) -> PreTrainedTokenizer:
         params = self._get_common_tokenizer_params()
  
@@ -86,17 +89,37 @@ class ModelTaskService(AsyncService):
             params["local_files_only"] = True
 
         return params
-    
-    def _get_model_device(self, model: PreTrainedModel) -> torch.device:
-        return next(model.parameters()).device
 
-    @abstractmethod
-    def _get_model_class(self) -> Type[PreTrainedModel]:
-        pass
-
-    @abstractmethod
     def _get_tokenizer_class(self) -> Type[PreTrainedTokenizer]:
-        pass
+        raise NotImplementedError("Tokenizer class loader not implemented.")
+
+    def _load_pretrained_processor(self, extra_params: Optional[Dict[str, Any]] = None) -> ProcessorMixin:
+        params = self._get_common_processor_params()
+ 
+        if extra_params:
+            params.update(extra_params)
+
+        return self._get_processor_class().from_pretrained(self.config.model, **params)
+
+    def _get_common_processor_params(self) -> Dict[str, Any]:
+        params: Dict[str, Any] = {}
+
+        if self.config.revision:
+            params["revision"] = self.config.revision
+
+        if self.config.cache_dir:
+            params["cache_dir"] = self.config.cache_dir
+
+        if self.config.local_files_only:
+            params["local_files_only"] = True
+
+        return params
+
+    def _get_processor_class(self) -> Type[ProcessorMixin]:
+        raise NotImplementedError("Processor class loader not implemented.")
+
+    def _get_model_device(self, model: PreTrainedModel) -> torch.device:
+        return next(model.parameters()).device
 
 def register_model_task_service(type: ModelTaskType):
     def decorator(cls: Type[ModelTaskService]) -> Type[ModelTaskService]:

{model_compose-0.3.9 → model_compose-0.3.10}/src/mindor/core/component/services/model/tasks/__init__.py

@@ -3,3 +3,4 @@ from .summarization import *
 from .translation import *
 from .text_classification import *
 from .text_embedding import *
+from .image_to_text import *

model_compose-0.3.10/src/mindor/core/component/services/model/tasks/image_to_text.py

@@ -0,0 +1,155 @@
+from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, Protocol, Any
+from mindor.dsl.schema.component import ModelComponentConfig, ImageToTextModelArchitecture
+from mindor.dsl.schema.action import ModelActionConfig, ImageToTextModelActionConfig
+from mindor.core.utils.streamer import AsyncStreamer
+from mindor.core.logger import logging
+from ..base import ModelTaskService, ModelTaskType, register_model_task_service
+from ..base import ComponentActionContext
+from transformers import PreTrainedModel, PreTrainedTokenizer, ProcessorMixin, GenerationMixin, TextIteratorStreamer
+from PIL import Image as PILImage
+from threading import Thread
+from torch import Tensor
+import torch, asyncio
+
+class WithTokenizer(Protocol):
+    tokenizer: PreTrainedTokenizer
+
+class ImageToTextTaskAction:
+    def __init__(self, config: ImageToTextModelActionConfig, model: PreTrainedModel, processor: ProcessorMixin, device: torch.device):
+        self.config: ImageToTextModelActionConfig = config
+        self.model: Union[PreTrainedModel, GenerationMixin] = model
+        self.processor: Union[ProcessorMixin, WithTokenizer] = processor
+        self.device: torch.device = device
+
+    async def run(self, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
+        image: Union[PILImage.Image, List[PILImage.Image]] = await context.render_image(self.config.image)
+        prompt: Optional[Union[str, List[str]]] = await context.render_variable(self.config.prompt)
+
+        batch_size = await context.render_variable(self.config.params.batch_size)
+        stream     = await context.render_variable(self.config.stream)
+
+        is_single_input: bool = True if not isinstance(images, list) else False
+        images: List[PILImage.Image] = [ image ] if is_single_input else image
+        prompts: Optional[List[str]] = [ prompt ] if is_single_input else prompt
+        results = []
+
+        if stream and (batch_size != 1 or len(images) != 1):
+            raise ValueError("Streaming mode only supports a single input image with batch size of 1.")
+
+        streamer = TextIteratorStreamer(self.processor.tokenizer, skip_prompt=True, skip_special_tokens=True) if stream else None
+        for index in range(0, len(images), batch_size):
+            batch_images = images[index:index + batch_size]
+            batch_prompts = prompts[index:index + batch_size] if prompts else None
+            
+            inputs: Tensor = self.processor(images=batch_images, text=batch_prompts, return_tensors="pt")
+            inputs = inputs.to(self.device)
+
+            def _generate():
+                with torch.inference_mode():
+                    outputs = self.model.generate(
+                        **inputs,
+                        streamer=streamer
+                    )
+
+                outputs = self.processor.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+                results.extend(outputs)
+
+            if stream:
+                thread = Thread(target=_generate)
+                thread.start()
+            else:
+                _generate()
+
+        if stream:
+            async def _stream_generator():
+                async for chunk in AsyncStreamer(streamer, loop):
+                    context.register_source("result", chunk)
+                    yield (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else chunk
+
+            return _stream_generator()
+        else:
+            result = results[0] if is_single_input else results
+            context.register_source("result", result)
+
+            return (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else result
+
+@register_model_task_service(ModelTaskType.IMAGE_TO_TEXT)
+class ImageToTextTaskService(ModelTaskService):
+    def __init__(self, id: str, config: ModelComponentConfig, daemon: bool):
+        super().__init__(id, config, daemon)
+
+        self.model: Optional[Union[PreTrainedModel, GenerationMixin]] = None
+        self.processor: Optional[ProcessorMixin] = None
+        self.device: Optional[torch.device] = None
+
+    async def _serve(self) -> None:
+        try:
+            self.model = self._load_pretrained_model()
+            self.processor = self._load_pretrained_processor()
+            self.device = self._get_model_device(self.model)
+            logging.info(f"Model and processor loaded successfully on device '{self.device}': {self.config.model}")
+        except Exception as e:
+            logging.error(f"Failed to load model '{self.config.model}': {e}")
+            raise
+
+    async def _shutdown(self) -> None:
+        self.model = None
+        self.processor = None
+        self.device = None
+
+    async def _run(self, action: ModelActionConfig, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
+        return await ImageToTextTaskAction(action, self.model, self.processor, self.device).run(context, loop)
+
+    def _get_model_class(self) -> Type[PreTrainedModel]:
+        if self.config.architecture == ImageToTextModelArchitecture.BLIP:
+            from transformers import BlipForConditionalGeneration
+            return BlipForConditionalGeneration
+
+        if self.config.architecture == ImageToTextModelArchitecture.BLIP2:
+            from transformers import Blip2ForConditionalGeneration
+            return Blip2ForConditionalGeneration
+
+        if self.config.architecture == ImageToTextModelArchitecture.GIT:
+            from transformers import GitForCausalLM
+            return GitForCausalLM
+
+        if self.config.architecture == ImageToTextModelArchitecture.PIX2STRUCT:
+            from transformers import Pix2StructForConditionalGeneration
+            return Pix2StructForConditionalGeneration
+
+        if self.config.architecture == ImageToTextModelArchitecture.DONUT:
+            from transformers import VisionEncoderDecoderModel # Donut uses this
+            return VisionEncoderDecoderModel
+
+        if self.config.architecture == ImageToTextModelArchitecture.KOSMOS2:
+            from transformers import Kosmos2ForConditionalGeneration
+            return Kosmos2ForConditionalGeneration
+        
+        raise ValueError(f"Unknown architecture: {self.config.architecture}")
+
+    def _get_processor_class(self) -> Type[ProcessorMixin]:
+        if self.config.architecture == ImageToTextModelArchitecture.BLIP:
+            from transformers import BlipProcessor
+            return BlipProcessor
+
+        if self.config.architecture == ImageToTextModelArchitecture.BLIP2:
+            from transformers import Blip2Processor
+            return Blip2Processor
+
+        if self.config.architecture == ImageToTextModelArchitecture.GIT:
+            from transformers import GitProcessor
+            return GitProcessor
+
+        if self.config.architecture == ImageToTextModelArchitecture.PIX2STRUCT:
+            from transformers import Pix2StructProcessor
+            return Pix2StructProcessor
+
+        if self.config.architecture == ImageToTextModelArchitecture.DONUT:
+            from transformers import DonutProcessor
+            return DonutProcessor
+
+        if self.config.architecture == ImageToTextModelArchitecture.KOSMOS2:
+            from transformers import Kosmos2Processor
+            return Kosmos2Processor
+
+        raise ValueError(f"Unknown architecture: {self.config.architecture}")

{model_compose-0.3.9 → model_compose-0.3.10}/src/mindor/core/component/services/model/tasks/summarization.py

@@ -11,7 +11,7 @@ from torch import Tensor
 import torch, asyncio
 
 class SummarizationTaskAction:
-    def __init__(self, config: SummarizationModelActionConfig, model: Union[PreTrainedModel, GenerationMixin], tokenizer: PreTrainedTokenizer, device: torch.device):
+    def __init__(self, config: SummarizationModelActionConfig, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, device: torch.device):
         self.config: SummarizationModelActionConfig = config
         self.model: Union[PreTrainedModel, GenerationMixin] = model
         self.tokenizer: PreTrainedTokenizer = tokenizer
@@ -32,7 +32,8 @@ class SummarizationTaskAction:
         batch_size        = await context.render_variable(self.config.params.batch_size)
         stream            = await context.render_variable(self.config.stream)
 
-        texts: List[str] = [ text ] if isinstance(text, str) else text
+        is_single_input: bool = True if not isinstance(text, list) else False
+        texts: List[str] = [ text ] if is_single_input else text
         results = []
 
         if stream and (batch_size != 1 or len(texts) != 1):
@@ -45,7 +46,7 @@ class SummarizationTaskAction:
             inputs = { k: v.to(self.device) for k, v in inputs.items() }
 
             def _generate():
-                with torch.no_grad():
+                with torch.inference_mode():
                     outputs = self.model.generate(
                         **inputs,
                         max_length=max_output_length,
@@ -61,8 +62,9 @@ class SummarizationTaskAction:
                         streamer=streamer
                     )
 
-                outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
-                results.extend(outputs)
+                if not stream:
+                    outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+                    results.extend(outputs)
 
             if stream:
                 thread = Thread(target=_generate)
@@ -70,15 +72,15 @@ class SummarizationTaskAction:
             else:
                 _generate()
 
-        if stream:        
+        if stream:
             async def _stream_generator():
-                async for result in AsyncStreamer(streamer, loop):
-                    context.register_source("result", result)
-                    yield (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else result
+                async for chunk in AsyncStreamer(streamer, loop):
+                    context.register_source("result", chunk)
+                    yield (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else chunk
 
             return _stream_generator()
         else:
-            result = results if len(results) > 1 else results[0]
+            result = results[0] if is_single_input else results
             context.register_source("result", result)
 
             return (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else result

model_compose-0.3.10/src/mindor/core/component/services/model/tasks/text_classification.py

@@ -0,0 +1,122 @@
+from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, Any
+from mindor.dsl.schema.component import ModelComponentConfig
+from mindor.dsl.schema.action import ModelActionConfig, TextClassificationModelActionConfig
+from mindor.core.utils.streamer import AsyncStreamer
+from mindor.core.logger import logging
+from ..base import ModelTaskService, ModelTaskType, register_model_task_service
+from ..base import ComponentActionContext
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, PreTrainedModel, PreTrainedTokenizer
+from transformers.modeling_outputs import SequenceClassifierOutput
+import torch.nn.functional as F
+from torch import Tensor
+import torch, asyncio
+
+class TextClassificationTaskAction:
+    def __init__(self, config: TextClassificationModelActionConfig, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, device: torch.device):
+        self.config: TextClassificationModelActionConfig = config
+        self.model: PreTrainedModel = model
+        self.tokenizer: PreTrainedTokenizer = tokenizer
+        self.device: torch.device = device
+
+    async def run(self, context: ComponentActionContext, labels: Optional[List[str]]) -> Any:
+        text: Union[str, List[str]] = await context.render_variable(self.config.text)
+
+        return_probabilities = await context.render_variable(self.config.params.return_probabilities)
+        batch_size           = await context.render_variable(self.config.params.batch_size)
+        stream               = await context.render_variable(self.config.stream)
+
+        is_single_input: bool = True if not isinstance(text, list) else False
+        has_iterable_result: bool = context.has_reference("result[]", self.config.output)
+        texts: List[str] = [ text ] if is_single_input else text
+        results = []
+
+        async def _predict():
+            for index in range(0, len(texts), batch_size):
+                batch_texts = texts[index:index + batch_size]
+                inputs: Dict[str, Tensor] = self.tokenizer(batch_texts, return_tensors="pt", padding=True, truncation=True)
+                inputs = { k: v.to(self.device) for k, v in inputs.items() }
+
+                with torch.inference_mode():
+                    outputs: SequenceClassifierOutput = self.model(**inputs)
+                    logits = outputs.logits  # shape: (batch_size, num_classes)
+                    predictions = []
+
+                    if return_probabilities:
+                        probs = F.softmax(logits, dim=-1).cpu()
+                        for prob in probs:
+                            predicted_index = torch.argmax(prob).item()
+                            predictions.append({
+                                "label": labels[predicted_index] if labels else predicted_index,
+                                "probabilities": prob.tolist()
+                            })
+                    else:
+                        predicted_indices = torch.argmax(logits, dim=-1).tolist()
+                        for predicted_index in predicted_indices:
+                            predictions.append(labels[predicted_index] if labels else predicted_index)
+
+                if self.config.output and has_iterable_result:
+                    rendered_outputs = []
+                    for prediction in predictions:
+                        context.register_source("result[]", prediction)
+                        output = await context.render_variable(self.config.output, ignore_files=True)
+                        rendered_outputs.append(output)
+                    yield rendered_outputs
+                else:
+                    yield predictions
+
+        if stream:
+            async def _stream_generator():
+                async for predictions in _predict():
+                    if not has_iterable_result:
+                        for prediction in predictions:
+                            context.register_source("result", prediction)
+                            yield (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else result
+                    else:
+                        for prediction in predictions:
+                            yield prediction
+
+            return _stream_generator()
+        else:
+            async for predictions in _predict():
+                results.extend(predictions)
+
+            if not has_iterable_result:
+                result = results[0] if is_single_input else results
+                context.register_source("result", result)
+
+                return (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else result
+            else:
+                return results
+
+@register_model_task_service(ModelTaskType.TEXT_CLASSIFICATION)
+class TextClassificationTaskService(ModelTaskService):
+    def __init__(self, id: str, config: ModelComponentConfig, daemon: bool):
+        super().__init__(id, config, daemon)
+
+        self.model: Optional[PreTrainedModel] = None
+        self.tokenizer: Optional[PreTrainedTokenizer] = None
+        self.device: Optional[torch.device] = None
+
+    async def _serve(self) -> None:
+        try:
+            self.model = self._load_pretrained_model()
+            self.tokenizer = self._load_pretrained_tokenizer()
+            self.device = self._get_model_device(self.model)
+            logging.info(f"Model and tokenizer loaded successfully on device '{self.device}': {self.config.model}")
+        except Exception as e:
+            logging.error(f"Failed to load model '{self.config.model}': {e}")
+            raise
+
+    async def _shutdown(self) -> None:
+        self.model = None
+        self.tokenizer = None
+        self.device = None
+
+    async def _run(self, action: ModelActionConfig, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
+        return await TextClassificationTaskAction(action, self.model, self.tokenizer, self.device).run(context, self.config.labels)
+
+    def _get_model_class(self) -> Type[PreTrainedModel]:
+        return AutoModelForSequenceClassification
+
+    def _get_tokenizer_class(self) -> Type[PreTrainedTokenizer]:
+        return AutoTokenizer

{model_compose-0.3.9 → model_compose-0.3.10}/src/mindor/core/component/services/model/tasks/text_embedding.py

@@ -1,6 +1,7 @@
 from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, Any
 from mindor.dsl.schema.component import ModelComponentConfig
 from mindor.dsl.schema.action import ModelActionConfig, TextEmbeddingModelActionConfig
+from mindor.core.utils.streamer import AsyncStreamer
 from mindor.core.logger import logging
 from ..base import ModelTaskService, ModelTaskType, register_model_task_service
 from ..base import ComponentActionContext
@@ -23,31 +24,64 @@ class TextEmbeddingTaskAction:
         pooling          = await context.render_variable(self.config.params.pooling)
         normalize        = await context.render_variable(self.config.params.normalize)
         batch_size       = await context.render_variable(self.config.params.batch_size)
+        stream           = await context.render_variable(self.config.stream)
 
-        texts: List[str] = [ text ] if isinstance(text, str) else text
+        is_single_input: bool = True if not isinstance(text, list) else False
+        has_iterable_result: bool = context.has_reference("result[]", self.config.output)
+        texts: List[str] = [ text ] if is_single_input else text
         results = []
 
-        for index in range(0, len(texts), batch_size):
-            batch_texts = texts[index:index + batch_size]
-            inputs: Dict[str, Tensor] = self.tokenizer(batch_texts, return_tensors="pt", max_length=max_input_length, padding=True, truncation=True)
-            inputs = { k: v.to(self.device) for k, v in inputs.items() }
-
-            with torch.no_grad():
-                outputs: BaseModelOutput = self.model(**inputs)
-                last_hidden_state = outputs.last_hidden_state  # (batch_size, seq_len, hidden_size)
-
-            attention_mask = inputs.get("attention_mask", None)
-            embeddings = self._pool_hidden_state(last_hidden_state, attention_mask, pooling)
-
-            if normalize:
-                embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
-
-            results.extend(embeddings.cpu().tolist())
-
-        result = results if len(results) > 1 else results[0]
-        context.register_source("result", result)
-
-        return (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else result
+        async def _embed():
+            for index in range(0, len(texts), batch_size):
+                batch_texts = texts[index:index + batch_size]
+                inputs: Dict[str, Tensor] = self.tokenizer(batch_texts, return_tensors="pt", max_length=max_input_length, padding=True, truncation=True)
+                inputs = { k: v.to(self.device) for k, v in inputs.items() }
+
+                with torch.inference_mode():
+                    outputs: BaseModelOutput = self.model(**inputs)
+                    last_hidden_state = outputs.last_hidden_state  # (batch_size, seq_len, hidden_size)
+
+                attention_mask = inputs.get("attention_mask", None)
+                embeddings = self._pool_hidden_state(last_hidden_state, attention_mask, pooling)
+
+                if normalize:
+                    embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1, eps=1e-12)
+
+                embeddings = embeddings.cpu().tolist()
+
+                if self.config.output and has_iterable_result:
+                    rendered_outputs = []
+                    for embedding in embeddings:
+                        context.register_source("result[]", embedding)
+                        output = await context.render_variable(self.config.output, ignore_files=True)
+                        rendered_outputs.append(output)
+                    yield rendered_outputs
+                else:
+                    yield embeddings
+
+        if stream:
+            async def _stream_generator():
+                async for embeddings in _embed():
+                    if not has_iterable_result:
+                        for embedding in embeddings:
+                            context.register_source("result", embedding)
+                            yield (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else result
+                    else:
+                        for embedding in embeddings:
+                            yield embedding
+
+            return _stream_generator()
+        else:
+            async for embeddings in _embed():
+                results.extend(embeddings)
+
+            if not has_iterable_result:
+                result = results[0] if is_single_input else results
+                context.register_source("result", result)
+
+                return (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else result
+            else:
+                return results
 
     def _pool_hidden_state(self, last_hidden_state: Tensor, attention_mask: Optional[Tensor], pooling: str) -> Tensor:
         if pooling == "mean":

{model_compose-0.3.9 → model_compose-0.3.10}/src/mindor/core/component/services/model/tasks/text_generation.py

@@ -11,7 +11,7 @@ from torch import Tensor
 import torch, asyncio
 
 class TextGenerationTaskAction:
-    def __init__(self, config: TextGenerationModelActionConfig, model: Union[PreTrainedModel, GenerationMixin], tokenizer: PreTrainedTokenizer, device: torch.device):
+    def __init__(self, config: TextGenerationModelActionConfig, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, device: torch.device):
         self.config: TextGenerationModelActionConfig = config
         self.model: Union[PreTrainedModel, GenerationMixin] = model
         self.tokenizer: PreTrainedTokenizer = tokenizer
@@ -28,7 +28,8 @@ class TextGenerationTaskAction:
         batch_size           = await context.render_variable(self.config.params.batch_size)
         stream               = await context.render_variable(self.config.stream)
 
-        prompts: List[str] = [ prompt ] if isinstance(prompt, str) else prompt
+        is_single_input: bool = True if not isinstance(prompt, list) else False
+        prompts: List[str] = [ prompt ] if is_single_input else prompt
         results = []
 
         if stream and (batch_size != 1 or len(prompts) != 1):
@@ -41,7 +42,7 @@ class TextGenerationTaskAction:
             inputs = { k: v.to(self.device) for k, v in inputs.items() }
 
             def _generate():
-                with torch.no_grad():
+                with torch.inference_mode():
                     outputs = self.model.generate(
                         **inputs,
                         max_new_tokens=max_output_length,
@@ -53,8 +54,9 @@ class TextGenerationTaskAction:
                         streamer=streamer
                     )
 
-                outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
-                results.extend(outputs)
+                if not stream:
+                    outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+                    results.extend(outputs)
 
             if stream:
                 thread = Thread(target=_generate)
@@ -62,15 +64,15 @@ class TextGenerationTaskAction:
             else:
                 _generate()
 
-        if stream:        
+        if stream:
             async def _stream_generator():
-                async for result in AsyncStreamer(streamer, loop):
-                    context.register_source("result", result)
-                    yield (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else result
+                async for chunk in AsyncStreamer(streamer, loop):
+                    context.register_source("result", chunk)
+                    yield (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else chunk
 
             return _stream_generator()
         else:
-            result = results if len(results) > 1 else results[0] 
+            result = results[0] if is_single_input else results
             context.register_source("result", result)
 
             return (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else result