model-compose 0.3.5__tar.gz → 0.3.7__tar.gz

{model_compose-0.3.5/src/model_compose.egg-info → model_compose-0.3.7}/PKG-INFO

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

{model_compose-0.3.5 → model_compose-0.3.7}/pyproject.toml

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

{model_compose-0.3.5 → model_compose-0.3.7}/src/mindor/cli/compose.py

@@ -4,14 +4,14 @@ from pathlib import Path
 import asyncio
 
 from mindor.dsl.loader import load_compose_config
-from mindor.core.runtime.env import load_env_files
+from mindor.core.runtime.env import load_env_files, merge_env_data
 from mindor.core.compose import *
 
 @click.group()
 @click.option(
     "--file", "-f", "config_files", multiple=True,
     type=click.Path(exists=True, dir_okay=False, path_type=Path),
-    help="Compose configuration files"
+    help="Compose configuration files."
 )
 @click.pass_context
 def compose_command(ctx: click.Context, config_files: List[Path]) -> None:
@@ -44,6 +44,7 @@ def up_command(
     async def _async_command():
         try:
             env = load_env_files(".", env_files or [])
+            env = merge_env_data(env, env_data)
             config = load_compose_config(".", config_files, env)
             await launch_services(config, detach, verbose)
         except Exception as e:
@@ -73,6 +74,7 @@ def down_command(
     async def _async_command():
         try:
             env = load_env_files(".", env_files or [])
+            env = merge_env_data(env, env_data)
             config = load_compose_config(".", config_files, env)
             await terminate_services(config, verbose)
         except Exception as e:
@@ -102,6 +104,7 @@ def start_command(
     async def _async_command():
         try:
             env = load_env_files(".", env_files or [])
+            env = merge_env_data(env, env_data)
             config = load_compose_config(".", config_files, env)
             await start_services(config, verbose)
         except Exception as e:
@@ -131,6 +134,7 @@ def stop_command(
     async def _async_command():
         try:
             env = load_env_files(".", env_files or [])
+            env = merge_env_data(env, env_data)
             config = load_compose_config(".", config_files, env)
             await stop_services(config, verbose)
         except Exception as e:
@@ -176,6 +180,7 @@ def run_command(
     async def _async_command():
         try:
             env = load_env_files(".", env_files or [])
+            env = merge_env_data(env, env_data)
             config = load_compose_config(".", config_files, env)
             input = json.loads(input_json) if input_json else {}
             state = await run_workflow(config, workflow_id, input, output_path, verbose)

{model_compose-0.3.5 → model_compose-0.3.7}/src/mindor/core/component/base.py

@@ -57,6 +57,12 @@ class ComponentService(AsyncService):
         if self.config.max_concurrent_count > 0:
             self.queue = WorkQueue(self.config.max_concurrent_count, self._run)
 
+    async def install(self) -> None:
+        await self._install()
+
+    async def build(self) -> None:
+        await self._build()
+
     async def run(self, action_id: Union[str, None], run_id: str, input: Dict[str, Any]) -> Dict[str, Any]:
         _, action = ActionResolver(self.config.actions).resolve(action_id)
         context = ComponentActionContext(run_id, input)
@@ -66,6 +72,12 @@ class ComponentService(AsyncService):
 
         return await self._run(action, context)
 
+    async def _install(self) -> None:
+        pass
+
+    async def _build(self) -> None:
+        pass
+
     async def _start(self) -> None:
         if self.queue:
             await self.queue.start()

{model_compose-0.3.5 → model_compose-0.3.7}/src/mindor/core/component/context.py

@@ -19,17 +19,23 @@ class ComponentActionContext:
     async def _resolve_source(self, key: str, index: Optional[int]) -> Any:
         if key in self.sources:
             return self.sources[key][index] if index is not None else self.sources[key]
+
         if key == "input":
             return self.input
+
         if key == "context":
             return self.context
+
         if key.startswith("gateway:"):
             return self._resolve_gateway(key)
+
         raise KeyError(f"Unknown source: {key}")
 
     def _resolve_gateway(self, key: str) -> Any:
         _, port = key.split(":")
         gateway = find_gateway_by_port(int(port)) if port else None
+
         if gateway:
             return gateway.get_context()
+
         return None

model_compose-0.3.7/src/mindor/core/component/services/http_server.py

@@ -0,0 +1,133 @@
+from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, Callable, Any
+from abc import ABC, abstractmethod
+from mindor.dsl.schema.component import HttpServerComponentConfig
+from mindor.dsl.schema.action import ActionConfig, HttpServerActionConfig, HttpServerCompletionType, HttpServerCompletionConfig
+from mindor.core.listener import HttpCallbackListener
+from mindor.core.utils.http_client import HttpClient
+from mindor.core.utils.http_status import is_status_code_matched
+from mindor.core.utils.time import parse_duration
+from mindor.core.utils.shell import run_command_streaming
+from ..base import ComponentService, ComponentType, ComponentGlobalConfigs, register_component
+from ..context import ComponentActionContext
+from datetime import datetime, timezone
+import asyncio, os
+
+class HttpServerCompletion(ABC):
+    def __init__(self, config: HttpServerCompletionConfig):
+        self.config: HttpServerCompletionConfig = config
+
+    @abstractmethod
+    async def run(self, context: ComponentActionContext, client: HttpClient) -> Any:
+        pass
+
+class HttpServerPollingCompletion(HttpServerCompletion):
+    async def run(self, context: ComponentActionContext, client: HttpClient) -> Any:
+        path    = await context.render_variable(self.config.path)
+        method  = await context.render_variable(self.config.method)
+        params  = await context.render_variable(self.config.params)
+        body    = await context.render_variable(self.config.body)
+        headers = await context.render_variable(self.config.headers)
+
+        interval = parse_duration((await context.render_variable(self.config.interval)) or 5.0)
+        timeout  = parse_duration((await context.render_variable(self.config.timeout)) or 300.0)
+        deadline = datetime.now(timezone.utc) + timeout
+
+        await asyncio.sleep(interval.total_seconds())
+
+        while datetime.now(timezone.utc) < deadline:
+            response, status_code = await client.request(path or "", method, params, body, headers, raise_on_error=False)
+            context.register_source("result", response)
+
+            status = (await context.render_variable(self.config.status)) if self.config.status else None
+            if status:
+                if status in (self.config.success_when or []):
+                    return response
+                if status in (self.config.fail_when or []):
+                    raise RuntimeError(f"Polling failed: status '{status}' matched a failure condition.")
+            else: # use status code
+                if is_status_code_matched(status_code, self.config.success_when or []):
+                    return response
+                if is_status_code_matched(status_code, self.config.fail_when or []):
+                    raise RuntimeError(f"Polling failed: status code '{status_code}' matched a failure condition.")
+
+            await asyncio.sleep(interval.total_seconds())
+
+        raise TimeoutError(f"Polling timed out after {timeout}.")
+
+class HttpServerCallbackCompletion(HttpServerCompletion):
+    async def run(self, context: ComponentActionContext, client: HttpClient) -> Any:
+        callback_id = await context.render_variable(self.config.wait_for)
+        future: asyncio.Future = asyncio.get_running_loop().create_future()
+        HttpCallbackListener.register_pending_future(callback_id, future)
+
+        return await future
+
+class HttpServerAction:
+    def __init__(self, config: HttpServerActionConfig):
+        self.config: HttpServerActionConfig = config
+        self.completion: HttpServerCompletion = None
+
+        if self.config.completion:
+            self._configure_completion()
+
+    def _configure_completion(self) -> None:
+        if self.config.completion.type == HttpServerCompletionType.POLLING:
+            self.completion = HttpServerPollingCompletion(self.config.completion)
+            return
+        
+        if self.config.completion.type == HttpServerCompletionType.CALLBACK:
+            self.completion = HttpServerCallbackCompletion(self.config.completion)
+            return
+        
+        raise ValueError(f"Unsupported http completion type: {self.config.completion.type}")
+
+    async def run(self, context: ComponentActionContext, client: HttpClient) -> Any:
+        path    = await context.render_variable(self.config.path)
+        method  = await context.render_variable(self.config.method)
+        params  = await context.render_variable(self.config.params)
+        body    = await context.render_variable(self.config.body)
+        headers = await context.render_variable(self.config.headers)
+
+        response, result = await client.request(path or "", method, params, body, headers), None
+        context.register_source("response", response)
+
+        if self.completion:
+            result = await self.completion.run(context, client)
+            context.register_source("result", result)
+
+        return (await context.render_variable(self.config.output, ignore_files=True)) if self.config.output else (result or response)
+
+@register_component(ComponentType.HTTP_SERVER)
+class HttpServerComponent(ComponentService):
+    def __init__(self, id: str, config: HttpServerComponentConfig, global_configs: ComponentGlobalConfigs, daemon: bool):
+        super().__init__(id, config, global_configs, daemon)
+
+        self.client: Optional[HttpClient] = None
+
+    async def _install(self) -> None:
+        if self.config.commands.install:
+            await run_command_streaming(self.config.commands.install, self.config.working_dir, self.config.env)
+
+    async def _build(self) -> None:
+        if self.config.commands.build:
+            await run_command_streaming(self.config.commands.build, self.config.working_dir, self.config.env)
+
+    async def _start(self) -> None:
+        base_url = f"http://localhost:{self.config.port}" + (self.config.base_path or "")
+        self.client = HttpClient(base_url, self.config.headers)
+        await super()._start()
+
+    async def _stop(self) -> None:
+        await super()._stop()
+        await self.client.close()
+        self.client = None
+
+    async def _serve(self) -> None:
+        if self.config.commands.start:
+            await run_command_streaming(self.config.commands.start, self.config.working_dir, self.config.env, block=False)
+
+    async def _shutdown(self) -> None:
+        pass
+
+    async def _run(self, action: ActionConfig, context: ComponentActionContext) -> Any:
+        return await HttpServerAction(action).run(context, self.client)

{model_compose-0.3.5 → model_compose-0.3.7}/src/mindor/core/component/services/model/base.py

@@ -5,7 +5,6 @@ from mindor.dsl.schema.action import ModelActionConfig
 from mindor.core.services import AsyncService
 from ...context import ComponentActionContext
 from transformers import PreTrainedModel, PreTrainedTokenizer
-from threading import Thread
 import torch, asyncio
 
 class ModelTaskService(AsyncService):
@@ -16,13 +15,15 @@ class ModelTaskService(AsyncService):
         self.config: ModelComponentConfig = config
 
     async def run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
+        loop: asyncio.AbstractEventLoop = asyncio.get_running_loop()
+
         async def _run():
-            return await self._run(action, context)
+            return await self._run(action, context, loop)
 
         return await self.run_in_thread(_run)
 
     @abstractmethod
-    async def _run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
+    async def _run(self, action: ModelActionConfig, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
         pass
 
     def _load_pretrained_model(self, extra_params: Optional[Dict[str, Any]] = None) -> PreTrainedModel:

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

@@ -1,12 +1,14 @@
 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, SummarizationModelActionConfig
+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 AutoModelForSeq2SeqLM, AutoTokenizer, PreTrainedModel, PreTrainedTokenizer, GenerationMixin
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, PreTrainedModel, PreTrainedTokenizer, GenerationMixin, TextIteratorStreamer
+from threading import Thread
 from torch import Tensor
-import torch
+import torch, asyncio
 
 class SummarizationTaskAction:
     def __init__(self, config: SummarizationModelActionConfig, model: Union[PreTrainedModel, GenerationMixin], tokenizer: PreTrainedTokenizer, device: torch.device):
@@ -15,46 +17,71 @@ class SummarizationTaskAction:
         self.tokenizer: PreTrainedTokenizer = tokenizer
         self.device: torch.device = device
 
-    async def run(self, context: ComponentActionContext) -> Any:
+    async def run(self, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
         text: Union[str, List[str]] = await context.render_variable(self.config.text)
 
         max_input_length  = await context.render_variable(self.config.params.max_input_length)
         max_output_length = await context.render_variable(self.config.params.max_output_length)
         min_output_length = await context.render_variable(self.config.params.min_output_length)
         num_beams         = await context.render_variable(self.config.params.num_beams)
-        length_penalty    = await context.render_variable(self.config.params.length_penalty)
-        early_stopping    = await context.render_variable(self.config.params.early_stopping)
+        length_penalty    = await context.render_variable(self.config.params.length_penalty) if num_beams > 1 else None
+        early_stopping    = await context.render_variable(self.config.params.early_stopping) if num_beams > 1 else False
         do_sample         = await context.render_variable(self.config.params.do_sample)
+        top_k             = await context.render_variable(self.config.params.top_k) if do_sample else None
+        top_p             = await context.render_variable(self.config.params.top_p) if do_sample else None
         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
         results = []
 
+        if stream and (batch_size != 1 or len(texts) != 1):
+            raise ValueError("Streaming mode only supports a single input text with batch size of 1.")
+
+        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True) if stream else None
         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 = self.model.generate(
-                    **inputs,
-                    max_length=max_output_length,
-                    min_length=min_output_length,
-                    num_beams=num_beams,
-                    length_penalty=length_penalty,
-                    early_stopping=early_stopping,
-                    do_sample=do_sample,
-                    pad_token_id=getattr(self.tokenizer, "pad_token_id", None),
-                    eos_token_id=getattr(self.tokenizer, "eos_token_id", None)
-                )
-
-            outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
-            results.extend(outputs)
-        
-        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
+            def _generate():
+                with torch.no_grad():
+                    outputs = self.model.generate(
+                        **inputs,
+                        max_length=max_output_length,
+                        min_length=min_output_length,
+                        num_beams=num_beams,
+                        length_penalty=length_penalty,
+                        early_stopping=early_stopping,
+                        do_sample=do_sample,
+                        top_k=top_k,
+                        top_p=top_p,
+                        pad_token_id=getattr(self.tokenizer, "pad_token_id", None),
+                        eos_token_id=getattr(self.tokenizer, "eos_token_id", None),
+                        streamer=streamer
+                    )
+
+                outputs = self.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 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
+
+            return _stream_generator()
+        else:
+            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
 
 @register_model_task_service(ModelTaskType.SUMMARIZATION)
 class SummarizationTaskService(ModelTaskService):
@@ -80,8 +107,8 @@ class SummarizationTaskService(ModelTaskService):
         self.tokenizer = None
         self.device = None
 
-    async def _run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
-        return await SummarizationTaskAction(action, self.model, self.tokenizer, self.device).run(context)
+    async def _run(self, action: ModelActionConfig, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
+        return await SummarizationTaskAction(action, self.model, self.tokenizer, self.device).run(context, loop)
 
     def _get_model_class(self) -> Type[PreTrainedModel]:
         return AutoModelForSeq2SeqLM

{model_compose-0.3.5 → model_compose-0.3.7}/src/mindor/core/component/services/model/tasks/text_classification.py

@@ -8,7 +8,7 @@ from transformers import AutoModelForSequenceClassification, AutoTokenizer, PreT
 from transformers.modeling_outputs import SequenceClassifierOutput
 import torch.nn.functional as F
 from torch import Tensor
-import torch
+import torch, asyncio
 
 class TextClassificationTaskAction:
     def __init__(self, config: TextClassificationModelActionConfig, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, device: torch.device):
@@ -80,7 +80,7 @@ class TextClassificationTaskService(ModelTaskService):
         self.tokenizer = None
         self.device = None
 
-    async def _run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
+    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]:

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

@@ -7,7 +7,7 @@ from ..base import ComponentActionContext
 from transformers import AutoModel, AutoTokenizer, PreTrainedModel, PreTrainedTokenizer
 from transformers.modeling_outputs import BaseModelOutput
 from torch import Tensor
-import torch
+import torch, asyncio
 
 class TextEmbeddingTaskAction:
     def __init__(self, config: TextEmbeddingModelActionConfig, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, device: torch.device):
@@ -94,7 +94,7 @@ class TextEmbeddingTaskService(ModelTaskService):
         self.tokenizer = None
         self.device = None
 
-    async def _run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
+    async def _run(self, action: ModelActionConfig, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
         return await TextEmbeddingTaskAction(action, self.model, self.tokenizer, self.device).run(context)
 
     def _get_model_class(self) -> Type[PreTrainedModel]:

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

@@ -1,12 +1,14 @@
 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, TextGenerationModelActionConfig
+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 AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PreTrainedTokenizer, GenerationMixin
+from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, PreTrainedTokenizer, GenerationMixin, TextIteratorStreamer
+from threading import Thread
 from torch import Tensor
-import torch
+import torch, asyncio
 
 class TextGenerationTaskAction:
     def __init__(self, config: TextGenerationModelActionConfig, model: Union[PreTrainedModel, GenerationMixin], tokenizer: PreTrainedTokenizer, device: torch.device):
@@ -15,7 +17,7 @@ class TextGenerationTaskAction:
         self.tokenizer: PreTrainedTokenizer = tokenizer
         self.device: torch.device = device
 
-    async def run(self, context: ComponentActionContext) -> Any:
+    async def run(self, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
         prompt: Union[str, List[str]] = await context.render_variable(self.config.prompt)
 
         max_output_length    = await context.render_variable(self.config.params.max_output_length)
@@ -24,33 +26,54 @@ class TextGenerationTaskAction:
         top_k                = await context.render_variable(self.config.params.top_k)
         top_p                = await context.render_variable(self.config.params.top_p)
         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
         results = []
 
+        if stream and (batch_size != 1 or len(prompts) != 1):
+            raise ValueError("Streaming mode only supports a single input prompt with batch size of 1.")
+
+        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True) if stream else None
         for index in range(0, len(prompts), batch_size):
             batch_prompts = prompts[index:index + batch_size]
             inputs: Dict[str, Tensor] = self.tokenizer(batch_prompts, return_tensors="pt", padding=True, truncation=True)
             inputs = { k: v.to(self.device) for k, v in inputs.items() }
 
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    **inputs,
-                    max_new_tokens=max_output_length,
-                    num_return_sequences=num_return_sequences,
-                    temperature=temperature,
-                    top_k=top_k,
-                    top_p=top_p,
-                    do_sample=True
-                )
-
-            outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
-            results.extend(outputs)
-
-        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
+            def _generate():
+                with torch.no_grad():
+                    outputs = self.model.generate(
+                        **inputs,
+                        max_new_tokens=max_output_length,
+                        num_return_sequences=num_return_sequences,
+                        temperature=temperature,
+                        do_sample=True,
+                        top_k=top_k,
+                        top_p=top_p,
+                        streamer=streamer
+                    )
+
+                outputs = self.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 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
+
+            return _stream_generator()
+        else:
+            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
 
 @register_model_task_service(ModelTaskType.TEXT_GENERATION)
 class TextGenerationTaskService(ModelTaskService):
@@ -76,8 +99,8 @@ class TextGenerationTaskService(ModelTaskService):
         self.tokenizer = None
         self.device = None
 
-    async def _run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
-        return await TextGenerationTaskAction(action, self.model, self.tokenizer, self.device).run(context)
+    async def _run(self, action: ModelActionConfig, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
+        return await TextGenerationTaskAction(action, self.model, self.tokenizer, self.device).run(context, loop)
 
     def _get_model_class(self) -> Type[PreTrainedModel]:
         return AutoModelForCausalLM