model-compose 0.3.14__tar.gz → 0.3.16__tar.gz

{model_compose-0.3.14/src/model_compose.egg-info → model_compose-0.3.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: model-compose
-Version: 0.3.14
+Version: 0.3.16
 Summary: model-compose: Declarative AI Model and Workflow Orchestrator
 Author-email: Hanyeol Cho <hanyeol.cho@gmail.com>
 License-Expression: MIT
@@ -20,10 +20,6 @@ Requires-Dist: pyngrok
 Requires-Dist: ulid
 Requires-Dist: gradio
 Requires-Dist: Pillow
-Requires-Dist: transformers>=4.21.0
-Requires-Dist: torch
-Requires-Dist: sentencepiece
-Requires-Dist: accelerate
 Requires-Dist: docker
 Dynamic: license-file
 

{model_compose-0.3.14 → model_compose-0.3.16}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "model-compose"
-version = "0.3.14"
+version = "0.3.16"
 description = "model-compose: Declarative AI Model and Workflow Orchestrator"
 authors = [
     { name = "Hanyeol Cho", email = "hanyeol.cho@gmail.com" }
@@ -21,10 +21,6 @@ dependencies = [
     "ulid", 
     "gradio",
     "Pillow",
-    "transformers>=4.21.0",
-    "torch",
-    "sentencepiece",
-    "accelerate",
     "docker"
 ]
 

{model_compose-0.3.14 → model_compose-0.3.16}/src/mindor/core/component/base.py

@@ -7,6 +7,8 @@ from mindor.dsl.schema.gateway import GatewayConfig
 from mindor.dsl.schema.workflow import WorkflowConfig
 from mindor.core.services import AsyncService
 from mindor.core.utils.workqueue import WorkQueue
+from mindor.core.utils.package import extract_module_name, is_module_installed, install_package
+from mindor.core.logger import logging
 from .context import ComponentActionContext
 
 class ActionResolver:
@@ -59,11 +61,19 @@ class ComponentService(AsyncService):
             self.queue = WorkQueue(self.config.max_concurrent_count, self._run)
 
     async def setup(self) -> None:
+        dependencies = self._get_setup_requirements()
+        if dependencies:
+            await self._install_packages(dependencies)
+
         await self._setup()
 
     async def teardown(self) -> None:
         await self._teardown()
 
+    async def start(self, background: bool = False) -> None:
+        await super().start(background)
+        await self.wait_until_ready()
+
     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)
@@ -79,6 +89,9 @@ class ComponentService(AsyncService):
     async def _teardown(self) -> None:
         pass
 
+    async def _get_setup_requirements(self) -> Optional[List[str]]:
+        return None
+
     async def _start(self) -> None:
         if self.queue:
             await self.queue.start()
@@ -97,10 +110,20 @@ class ComponentService(AsyncService):
     async def _shutdown(self) -> None:
         pass
 
+    async def _is_ready(self) -> bool:
+        return True
+
     @abstractmethod
     async def _run(self, action: ActionConfig, context: ComponentActionContext) -> Any:
         pass
 
+    async def _install_packages(self, packages: List[str]) -> None:
+        for package_spec in packages:
+            module_name = extract_module_name(package_spec)
+            if not is_module_installed(module_name):
+                logging.info(f"Installing missing module: {package_spec}")
+                await install_package(package_spec)
+
 def register_component(type: ComponentType):
     def decorator(cls: Type[ComponentService]) -> Type[ComponentService]:
         ComponentRegistry[type] = cls

{model_compose-0.3.14 → model_compose-0.3.16}/src/mindor/core/component/services/http_server.py

@@ -161,10 +161,19 @@ class HttpServerComponent(ComponentService):
 
     async def _serve(self) -> None:
         if self.config.manage.scripts.start:
-            await run_command_streaming(self.config.manage.scripts.start, self.config.manage.working_dir, self.config.manage.env, block=False)
+            await run_command_streaming(self.config.manage.scripts.start, self.config.manage.working_dir, self.config.manage.env)
 
     async def _shutdown(self) -> None:
         pass
 
+    async def _is_ready(self) -> bool:
+        try:
+            _, writer = await asyncio.open_connection("localhost", self.config.port)
+            writer.close()
+            await writer.wait_closed()
+            return True
+        except (ConnectionRefusedError, OSError):
+            return False
+
     async def _run(self, action: ActionConfig, context: ComponentActionContext) -> Any:
         return await HttpServerAction(action).run(context, self.client)

{model_compose-0.3.14 → model_compose-0.3.16}/src/mindor/core/component/services/model/base.py

@@ -1,11 +1,17 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
 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 ModelComponentConfig, ModelTaskType, DeviceMode
+from mindor.dsl.schema.component import ModelComponentConfig, ModelTaskType, ModelSourceConfig, DeviceMode
 from mindor.dsl.schema.action import ModelActionConfig
 from mindor.core.services import AsyncService
 from ...context import ComponentActionContext
-from transformers import PreTrainedModel, PreTrainedTokenizer, ProcessorMixin
-import torch, asyncio
+import asyncio
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel, PreTrainedTokenizer, ProcessorMixin
+    import torch
 
 class ModelTaskService(AsyncService):
     def __init__(self, id: str, config: ModelComponentConfig, daemon: bool):
@@ -14,6 +20,9 @@ class ModelTaskService(AsyncService):
         self.id: str = id
         self.config: ModelComponentConfig = config
 
+    def get_setup_requirements(self) -> Optional[List[str]]:
+        return None
+
     async def run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
         loop: asyncio.AbstractEventLoop = asyncio.get_running_loop()
 
@@ -26,7 +35,7 @@ class ModelTaskService(AsyncService):
     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:
+    def _load_pretrained_model(self, extra_params: Optional[Dict[str, Any]] = None) -> Union[PreTrainedModel, torch.nn.Module]:
         params = self._get_common_model_params()
 
         if extra_params:
@@ -42,6 +51,10 @@ class ModelTaskService(AsyncService):
     def _get_common_model_params(self) -> Dict[str, Any]:
         params: Dict[str, Any] = {}
 
+        if isinstance(self.config.model, ModelSourceConfig):
+            if self.config.model.revision:
+                params["revision"] = self.config.model.revision
+
         if self.config.device_mode != DeviceMode.SINGLE:
             params["device_map"] = self.config.device_mode.value
     
@@ -51,9 +64,6 @@ class ModelTaskService(AsyncService):
         if self.config.low_cpu_mem_usage:
             params["low_cpu_mem_usage"] = True
 
-        if self.config.revision:
-            params["revision"] = self.config.revision
-
         if self.config.cache_dir:
             params["cache_dir"] = self.config.cache_dir
 
@@ -75,12 +85,13 @@ class ModelTaskService(AsyncService):
 
     def _get_common_tokenizer_params(self) -> Dict[str, Any]:
         params: Dict[str, Any] = {}
+
+        if isinstance(self.config.model, ModelSourceConfig):
+            if self.config.model.revision:
+                params["revision"] = self.config.model.revision
     
         if not self.config.fast_tokenizer:
             params["use_fast"] = False
-    
-        if self.config.revision:
-            params["revision"] = self.config.revision
 
         if self.config.cache_dir:
             params["cache_dir"] = self.config.cache_dir
@@ -104,9 +115,10 @@ class ModelTaskService(AsyncService):
     def _get_common_processor_params(self) -> Dict[str, Any]:
         params: Dict[str, Any] = {}
 
-        if self.config.revision:
-            params["revision"] = self.config.revision
-
+        if isinstance(self.config.model, ModelSourceConfig):
+            if self.config.model.revision:
+                params["revision"] = self.config.model.revision
+ 
         if self.config.cache_dir:
             params["cache_dir"] = self.config.cache_dir
 

model_compose-0.3.16/src/mindor/core/component/services/model/loader.py

@@ -0,0 +1,25 @@
+from typing import Type, Union, Literal, Optional, Dict, List, Tuple, Set, Annotated, AsyncIterator, Any
+from transformers import PreTrainedModel
+import torch
+    
+async def load_model(
+    provider: str,
+    model_id: str,
+    revision: Optional[str],
+    filename: Optional[str],
+    model_class: Type[PreTrainedModel],
+    *,
+    cache_dir: Optional[str] = None,
+    local_files_only: bool = False
+) -> Union[PreTrainedModel, torch.nn.Module]:
+    if provider in [ "huggingface.co", "huggingface.com", "huggingface" ]:
+        pass
+
+async def load_model_from_uri(
+    uri: str,
+    model_class: Type[PreTrainedModel],
+    *,
+    cache_dir: Optional[str] = None,
+    local_files_only: bool = False
+) -> Union[PreTrainedModel, torch.nn.Module]:
+    pass

{model_compose-0.3.14 → model_compose-0.3.16}/src/mindor/core/component/services/model/model.py

@@ -17,7 +17,7 @@ class ModelComponent(ComponentService):
     def __init__(self, id: str, config: ModelComponentConfig, global_configs: ComponentGlobalConfigs, daemon: bool):
         super().__init__(id, config, global_configs, daemon)
 
-        self.task_service = self._create_task_service(self.config.task)
+        self.task_service: ModelTaskService = self._create_task_service(self.config.task)
 
     def _create_task_service(self, type: ModelTaskType) -> ModelTaskService:
         try:
@@ -27,6 +27,17 @@ class ModelComponent(ComponentService):
         except KeyError:
             raise ValueError(f"Unsupported model task type: {type}")
 
+    def _get_setup_requirements(self) -> Optional[List[str]]:
+        task_dependencies = self.task_service.get_setup_requirements()
+
+        return [ 
+            "transformers>=4.21.0",
+            "torch",
+            "sentencepiece",
+            "accelerate",
+            *(task_dependencies or [])
+        ]
+
     async def _serve(self) -> None:
         await self.task_service.start()
 

{model_compose-0.3.14 → model_compose-0.3.16}/src/mindor/core/component/services/model/tasks/image_to_text.py

@@ -1,3 +1,6 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
 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
@@ -5,11 +8,14 @@ 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
+import asyncio
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel, PreTrainedTokenizer, ProcessorMixin, GenerationMixin, TextIteratorStreamer, StopStringCriteria
+    from torch import Tensor
+    import torch
 
 class WithTokenizer(Protocol):
     tokenizer: PreTrainedTokenizer
@@ -22,11 +28,25 @@ class ImageToTextTaskAction:
         self.device: torch.device = device
 
     async def run(self, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
+        import torch
+
         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)
+        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_return_sequences = await context.render_variable(self.config.params.num_return_sequences)
+        do_sample            = await context.render_variable(self.config.params.do_sample)
+        temperature          = await context.render_variable(self.config.params.temperature) if do_sample else None
+        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
+        num_beams            = await context.render_variable(self.config.params.num_beams)
+        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
+        stop_sequences       = await context.render_variable(self.config.params.stop_sequences)
+        batch_size           = await context.render_variable(self.config.params.batch_size)
+        stream               = await context.render_variable(self.config.stream)
 
         is_single_input: bool = bool(not isinstance(images, list))
         images: List[PILImage.Image] = [ image ] if is_single_input else image
@@ -37,17 +57,31 @@ class ImageToTextTaskAction:
             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
+        stopping_criteria = [ StopStringCriteria(self.processor.tokenizer, stop_sequences) ] if stop_sequences 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: Tensor = self.processor(images=batch_images, text=batch_prompts, max_length=max_input_length, return_tensors="pt")
             inputs = inputs.to(self.device)
 
             def _generate():
                 with torch.inference_mode():
                     outputs = self.model.generate(
                         **inputs,
+                        max_new_tokens=max_output_length,
+                        min_length=min_output_length,
+                        num_return_sequences=num_return_sequences,
+                        do_sample=do_sample,
+                        temperature=temperature,
+                        top_k=top_k,
+                        top_p=top_p,
+                        num_beams=num_beams,
+                        length_penalty=length_penalty,
+                        early_stopping=early_stopping,
+                        pad_token_id=getattr(self.processor.tokenizer, "pad_token_id", None),
+                        eos_token_id=getattr(self.processor.tokenizer, "eos_token_id", None),
+                        stopping_criteria=stopping_criteria,
                         streamer=streamer
                     )
 

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

@@ -1,3 +1,6 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
 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
@@ -5,10 +8,13 @@ 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, TextIteratorStreamer
-from threading import Thread
-from torch import Tensor
-import torch, asyncio
+import asyncio
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel, PreTrainedTokenizer, GenerationMixin, TextIteratorStreamer, StopStringCriteria
+    from threading import Thread
+    from torch import Tensor
+    import torch
 
 class SummarizationTaskAction:
     def __init__(self, config: SummarizationModelActionConfig, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, device: torch.device):
@@ -18,17 +24,21 @@ class SummarizationTaskAction:
         self.device: torch.device = device
 
     async def run(self, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
+        import torch
+
         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) 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)
+        temperature       = await context.render_variable(self.config.params.temperature) if do_sample else None
         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
+        num_beams         = await context.render_variable(self.config.params.num_beams)
+        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
+        stop_sequences    = await context.render_variable(self.config.params.stop_sequences)
         batch_size        = await context.render_variable(self.config.params.batch_size)
         stream            = await context.render_variable(self.config.stream)
 
@@ -40,6 +50,7 @@ class SummarizationTaskAction:
             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
+        stopping_criteria = [ StopStringCriteria(self.tokenizer, stop_sequences) ] if stop_sequences 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)
@@ -49,16 +60,18 @@ class SummarizationTaskAction:
                 with torch.inference_mode():
                     outputs = self.model.generate(
                         **inputs,
-                        max_length=max_output_length,
+                        max_new_tokens=max_output_length,
                         min_length=min_output_length,
-                        num_beams=num_beams,
-                        length_penalty=length_penalty,
-                        early_stopping=early_stopping,
                         do_sample=do_sample,
+                        temperature=temperature,
                         top_k=top_k,
                         top_p=top_p,
+                        num_beams=num_beams,
+                        length_penalty=length_penalty,
+                        early_stopping=early_stopping,
                         pad_token_id=getattr(self.tokenizer, "pad_token_id", None),
                         eos_token_id=getattr(self.tokenizer, "eos_token_id", None),
+                        stopping_criteria=stopping_criteria,
                         streamer=streamer
                     )
 
@@ -114,7 +127,9 @@ class SummarizationTaskService(ModelTaskService):
         return await SummarizationTaskAction(action, self.model, self.tokenizer, self.device).run(context, loop)
 
     def _get_model_class(self) -> Type[PreTrainedModel]:
+        from transformers import AutoModelForSeq2SeqLM
         return AutoModelForSeq2SeqLM
 
     def _get_tokenizer_class(self) -> Type[PreTrainedTokenizer]:
+        from transformers import AutoTokenizer
         return AutoTokenizer

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

@@ -1,14 +1,19 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
 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.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
+import asyncio
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel, PreTrainedTokenizer
+    from transformers.modeling_outputs import SequenceClassifierOutput
+    from torch import Tensor
+    import torch
 
 class TextClassificationTaskAction:
     def __init__(self, config: TextClassificationModelActionConfig, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, device: torch.device):
@@ -18,6 +23,8 @@ class TextClassificationTaskAction:
         self.device: torch.device = device
 
     async def run(self, context: ComponentActionContext, labels: Optional[List[str]]) -> Any:
+        import torch, torch.nn.functional as F
+
         text: Union[str, List[str]] = await context.render_variable(self.config.text)
 
         return_probabilities = await context.render_variable(self.config.params.return_probabilities)
@@ -115,7 +122,9 @@ class TextClassificationTaskService(ModelTaskService):
         return await TextClassificationTaskAction(action, self.model, self.tokenizer, self.device).run(context, self.config.labels)
 
     def _get_model_class(self) -> Type[PreTrainedModel]:
+        from transformers import AutoModelForSequenceClassification
         return AutoModelForSequenceClassification
 
     def _get_tokenizer_class(self) -> Type[PreTrainedTokenizer]:
+        from transformers import AutoTokenizer
         return AutoTokenizer

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

@@ -1,13 +1,19 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
 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.logger import logging
 from ..base import ModelTaskService, ModelTaskType, register_model_task_service
 from ..base import ComponentActionContext
-from transformers import AutoModel, AutoTokenizer, PreTrainedModel, PreTrainedTokenizer
-from transformers.modeling_outputs import BaseModelOutput
-from torch import Tensor
-import torch, asyncio
+import asyncio
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel, PreTrainedTokenizer
+    from transformers.modeling_outputs import BaseModelOutput
+    from torch import Tensor
+    import torch
 
 class TextEmbeddingTaskAction:
     def __init__(self, config: TextEmbeddingModelActionConfig, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, device: torch.device):
@@ -17,6 +23,8 @@ class TextEmbeddingTaskAction:
         self.device: torch.device = device
 
     async def run(self, context: ComponentActionContext) -> Any:
+        import torch, torch.nn.functional as F
+
         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)
@@ -44,7 +52,7 @@ class TextEmbeddingTaskAction:
                 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 = F.normalize(embeddings, p=2, dim=1, eps=1e-12)
 
                 embeddings = embeddings.cpu().tolist()
 
@@ -83,6 +91,8 @@ class TextEmbeddingTaskAction:
                 return results
 
     def _pool_hidden_state(self, last_hidden_state: Tensor, attention_mask: Optional[Tensor], pooling: str) -> Tensor:
+        import torch
+
         if pooling == "mean":
             if attention_mask is not None:
                 mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size())
@@ -131,7 +141,9 @@ class TextEmbeddingTaskService(ModelTaskService):
         return await TextEmbeddingTaskAction(action, self.model, self.tokenizer, self.device).run(context)
 
     def _get_model_class(self) -> Type[PreTrainedModel]:
+        from transformers import AutoModel
         return AutoModel
 
     def _get_tokenizer_class(self) -> Type[PreTrainedTokenizer]:
+        from transformers import AutoTokenizer
         return AutoTokenizer

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

@@ -1,3 +1,6 @@
+from __future__ import annotations
+from typing import TYPE_CHECKING
+
 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
@@ -5,10 +8,13 @@ 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, TextIteratorStreamer
-from threading import Thread
-from torch import Tensor
-import torch, asyncio
+import asyncio
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel, PreTrainedTokenizer, GenerationMixin, TextIteratorStreamer, StopStringCriteria
+    from threading import Thread
+    from torch import Tensor
+    import torch
 
 class TextGenerationTaskAction:
     def __init__(self, config: TextGenerationModelActionConfig, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, device: torch.device):
@@ -18,13 +24,21 @@ class TextGenerationTaskAction:
         self.device: torch.device = device
 
     async def run(self, context: ComponentActionContext, loop: asyncio.AbstractEventLoop) -> Any:
+        import torch
+
         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)
+        min_output_length    = await context.render_variable(self.config.params.min_output_length)
         num_return_sequences = await context.render_variable(self.config.params.num_return_sequences)
+        do_sample            = await context.render_variable(self.config.params.do_sample)
         temperature          = await context.render_variable(self.config.params.temperature)
         top_k                = await context.render_variable(self.config.params.top_k)
         top_p                = await context.render_variable(self.config.params.top_p)
+        num_beams            = await context.render_variable(self.config.params.num_beams)
+        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
+        stop_sequences       = await context.render_variable(self.config.params.stop_sequences)
         batch_size           = await context.render_variable(self.config.params.batch_size)
         stream               = await context.render_variable(self.config.stream)
 
@@ -36,6 +50,7 @@ class TextGenerationTaskAction:
             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
+        stopping_criteria = [ StopStringCriteria(self.tokenizer, stop_sequences) ] if stop_sequences 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)
@@ -46,11 +61,18 @@ class TextGenerationTaskAction:
                     outputs = self.model.generate(
                         **inputs,
                         max_new_tokens=max_output_length,
+                        min_length=min_output_length,
                         num_return_sequences=num_return_sequences,
+                        do_sample=do_sample,
                         temperature=temperature,
-                        do_sample=True,
                         top_k=top_k,
                         top_p=top_p,
+                        num_beams=num_beams,
+                        length_penalty=length_penalty,
+                        early_stopping=early_stopping,
+                        pad_token_id=getattr(self.tokenizer, "pad_token_id", None),
+                        eos_token_id=getattr(self.tokenizer, "eos_token_id", None),
+                        stopping_criteria=stopping_criteria,
                         streamer=streamer
                     )
 
@@ -106,7 +128,9 @@ class TextGenerationTaskService(ModelTaskService):
         return await TextGenerationTaskAction(action, self.model, self.tokenizer, self.device).run(context, loop)
 
     def _get_model_class(self) -> Type[PreTrainedModel]:
+        from transformers import AutoModelForCausalLM
         return AutoModelForCausalLM
 
     def _get_tokenizer_class(self) -> Type[PreTrainedTokenizer]:
+        from transformers import AutoTokenizer
         return AutoTokenizer