model-compose 0.3.2__tar.gz → 0.3.3__tar.gz

{model_compose-0.3.2/src/model_compose.egg-info → model_compose-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: model-compose
-Version: 0.3.2
+Version: 0.3.3
 Summary: model-compose: Declarative AI Model and Workflow Orchestrator
 Author-email: Hanyeol Cho <hanyeol.cho@gmail.com>
 License-Expression: MIT
@@ -248,7 +248,46 @@ Once enabled, the Web UI will be available at:
 ```
 http://localhost:8081
 ```
-> By default, the Web UI is powered by [Gradio](https://www.gradio.app), but support for other drivers may be added in the future.
+
+## 🎨 Custom Web UI
+You can fully customize the Web UI experience by specifying a different driver or serving your own frontend.
+
+By default, `model-compose` uses [Gradio](https://www.gradio.app) as the interactive UI. However, if you prefer to use your own static frontend (e.g., a custom React/Vite app), you can switch to the `static` driver.
+
+Here’s how you can do it:
+
+```
+controller:
+  type: http-server
+  port: 8080
+  webui:
+    driver: static
+    static_dir: webui
+    port: 8081
+```
+Your frontend should be a prebuilt static site (e.g., using `vite build`, `next export`, or `react-scripts` build) and placed in the specified `static_dir`.
+Make sure `index.html` exists in that directory.
+
+#### 📁 Example Directory Structure
+```
+project/
+├── model-compose.yml
+├── webui/
+│   ├── index.html
+│   ├── assets/
+│   └── ...
+```
+
+#### 🔧 Notes
+* Ensure the `static_dir` path is relative to the project root or an absolute path.
+* You can use environment variables inside `model-compose.yml` to make this path configurable.
+
+Support for additional drivers (e.g., `dynamic`) may be added in future versions.
+
+Once configured, the Web UI will be available at:
+```
+http://localhost:8081
+```
 
 ---
 ## 🛰 MCP Server Support

{model_compose-0.3.2 → model_compose-0.3.3}/README.md

@@ -219,7 +219,46 @@ Once enabled, the Web UI will be available at:
 ```
 http://localhost:8081
 ```
-> By default, the Web UI is powered by [Gradio](https://www.gradio.app), but support for other drivers may be added in the future.
+
+## 🎨 Custom Web UI
+You can fully customize the Web UI experience by specifying a different driver or serving your own frontend.
+
+By default, `model-compose` uses [Gradio](https://www.gradio.app) as the interactive UI. However, if you prefer to use your own static frontend (e.g., a custom React/Vite app), you can switch to the `static` driver.
+
+Here’s how you can do it:
+
+```
+controller:
+  type: http-server
+  port: 8080
+  webui:
+    driver: static
+    static_dir: webui
+    port: 8081
+```
+Your frontend should be a prebuilt static site (e.g., using `vite build`, `next export`, or `react-scripts` build) and placed in the specified `static_dir`.
+Make sure `index.html` exists in that directory.
+
+#### 📁 Example Directory Structure
+```
+project/
+├── model-compose.yml
+├── webui/
+│   ├── index.html
+│   ├── assets/
+│   └── ...
+```
+
+#### 🔧 Notes
+* Ensure the `static_dir` path is relative to the project root or an absolute path.
+* You can use environment variables inside `model-compose.yml` to make this path configurable.
+
+Support for additional drivers (e.g., `dynamic`) may be added in future versions.
+
+Once configured, the Web UI will be available at:
+```
+http://localhost:8081
+```
 
 ---
 ## 🛰 MCP Server Support

{model_compose-0.3.2 → model_compose-0.3.3}/pyproject.toml

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

{model_compose-0.3.2 → model_compose-0.3.3}/src/mindor/cli/compose.py

@@ -11,7 +11,7 @@ from mindor.core.compose import *
 @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:
@@ -25,11 +25,11 @@ def compose_command(ctx: click.Context, config_files: List[Path]) -> None:
     "--env-file", "env_files", multiple=True,
     type=click.Path(exists=True, dir_okay=False, path_type=Path),
     required=False,
-    help="Path to a .env file containing environment variables.",
+    help="Path to a .env file containing environment variables."
 )
 @click.option(
     "--env", "-e", "env_data", multiple=True,
-    help="Environment variable in the form KEY=VALUE. Repeatable.",
+    help="Environment variable in the form KEY=VALUE. Repeatable."
 )
 @click.option("--verbose", "-v", is_flag=True, help="Enable verbose output.")
 @click.pass_context
@@ -55,11 +55,11 @@ def up_command(
     "--env-file", "env_files", multiple=True,
     type=click.Path(exists=True, dir_okay=False, path_type=Path),
     required=False,
-    help="Path to a .env file containing environment variables.",
+    help="Path to a .env file containing environment variables."
 )
 @click.option(
     "--env", "-e", "env_data", multiple=True,
-    help="Environment variable in the form KEY=VALUE. Repeatable.",
+    help="Environment variable in the form KEY=VALUE. Repeatable."
 )
 @click.option("--verbose", "-v", is_flag=True, help="Enable verbose output.")
 @click.pass_context
@@ -84,7 +84,7 @@ def down_command(
     "--env-file", "env_files", multiple=True,
     type=click.Path(exists=True, dir_okay=False, path_type=Path),
     required=False,
-    help="Path to a .env file containing environment variables.",
+    help="Path to a .env file containing environment variables."
 )
 @click.option(
     "--env", "-e", "env_data", multiple=True,
@@ -113,11 +113,11 @@ def start_command(
     "--env-file", "env_files", multiple=True,
     type=click.Path(exists=True, dir_okay=False, path_type=Path),
     required=False,
-    help="Path to a .env file containing environment variables.",
+    help="Path to a .env file containing environment variables."
 )
 @click.option(
     "--env", "-e", "env_data", multiple=True,
-    help="Environment variable in the form KEY=VALUE. Repeatable.",
+    help="Environment variable in the form KEY=VALUE. Repeatable."
 )
 @click.option("--verbose", "-v", is_flag=True, help="Enable verbose output.")
 @click.pass_context
@@ -149,11 +149,11 @@ def stop_command(
     "--env-file", "env_files", multiple=True,
     type=click.Path(exists=True, dir_okay=False, path_type=Path),
     required=False,
-    help="Path to a .env file containing environment variables.",
+    help="Path to a .env file containing environment variables."
 )
 @click.option(
     "--env", "-e", "env_data", multiple=True,
-    help="Environment variable in the form KEY=VALUE. Repeatable.",
+    help="Environment variable in the form KEY=VALUE. Repeatable."
 )
 @click.option("--verbose", "-v", is_flag=True, help="Enable verbose output.")
 @click.pass_context

{model_compose-0.3.2 → model_compose-0.3.3}/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.renderer import VariableRenderer
+from mindor.core.utils.renderers import VariableRenderer
 from mindor.core.gateway import find_gateway_by_port
 
 class ComponentActionContext:

model_compose-0.3.3/src/mindor/core/component/services/model/base.py

@@ -0,0 +1,102 @@
+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.action import ModelActionConfig
+from mindor.core.services import AsyncService
+from ...context import ComponentActionContext
+from transformers import PreTrainedModel, PreTrainedTokenizer
+import torch
+
+class ModelTaskService(AsyncService):
+    def __init__(self, id: str, config: ModelComponentConfig, daemon: bool):
+        super().__init__(daemon)
+
+        self.id: str = id
+        self.config: ModelComponentConfig = config
+
+    async def run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
+        return await self._run(action, context)
+
+    @abstractmethod
+    async def _run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
+        pass
+
+    def _load_pretrained_model(self, extra_params: Optional[Dict[str, Any]] = None) -> PreTrainedModel:
+        params = self._get_common_model_params()
+
+        if extra_params:
+            params.update(extra_params)
+
+        model = self._get_model_class().from_pretrained(self.config.model, **params)
+
+        if self.config.device_mode == DeviceMode.SINGLE:
+            model = model.to(torch.device(self.config.device))
+
+        return model
+
+    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
+    
+        if self.config.precision is not None:
+            params["torch_dtype"] = getattr(torch, self.config.precision.value)
+    
+        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
+
+        if self.config.local_files_only:
+            params["local_files_only"] = True
+
+        return params
+    
+    def _load_pretrained_tokenizer(self, extra_params: Optional[Dict[str, Any]] = None) -> PreTrainedTokenizer:
+        params = self._get_common_tokenizer_params()
+ 
+        if extra_params:
+            params.update(extra_params)
+
+        return self._get_tokenizer_class().from_pretrained(self.config.model, **params)
+
+    def _get_common_tokenizer_params(self) -> Dict[str, Any]:
+        params: Dict[str, Any] = {}
+    
+        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
+
+        if self.config.local_files_only:
+            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
+
+def register_model_task_service(type: ModelTaskType):
+    def decorator(cls: Type[ModelTaskService]) -> Type[ModelTaskService]:
+        ModelTaskServiceRegistry[type] = cls
+        return cls
+    return decorator
+
+ModelTaskServiceRegistry: Dict[ModelTaskType, Type[ModelTaskService]] = {}

{model_compose-0.3.2 → model_compose-0.3.3}/src/mindor/core/component/services/model/model.py

@@ -25,7 +25,7 @@ class ModelComponent(ComponentService):
                 from . import tasks
             return ModelTaskServiceRegistry[type](self.id, self.config, self.daemon)
         except KeyError:
-            raise ValueError(f"Unsupported component type: {type}")
+            raise ValueError(f"Unsupported model task type: {type}")
 
     async def _serve(self) -> None:
         await self.task_service.start()

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

@@ -5,13 +5,15 @@ 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 torch import Tensor
 import torch
 
 class SummarizationTaskAction:
-    def __init__(self, config: SummarizationModelActionConfig, model: Union[PreTrainedModel, GenerationMixin], tokenizer: PreTrainedTokenizer):
+    def __init__(self, config: SummarizationModelActionConfig, model: Union[PreTrainedModel, GenerationMixin], tokenizer: PreTrainedTokenizer, device: torch.device):
         self.config: SummarizationModelActionConfig = config
         self.model: Union[PreTrainedModel, GenerationMixin] = model
         self.tokenizer: PreTrainedTokenizer = tokenizer
+        self.device: torch.device = device
 
     async def run(self, context: ComponentActionContext) -> Any:
         text: Union[str, List[str]] = await context.render_variable(self.config.text)
@@ -30,7 +32,8 @@ class SummarizationTaskAction:
 
         for index in range(0, len(texts), batch_size):
             batch_texts = texts[index:index + batch_size]
-            inputs = self.tokenizer(batch_texts, return_tensors="pt", max_length=max_input_length, padding=True, truncation=True).to(self.model.device)
+            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(
@@ -60,12 +63,14 @@ class SummarizationTaskService(ModelTaskService):
 
         self.model: Optional[Union[PreTrainedModel, GenerationMixin]] = None
         self.tokenizer: Optional[PreTrainedTokenizer] = None
+        self.device: Optional[torch.device] = None
 
     async def _serve(self) -> None:
         try:
-            self.model = AutoModelForSeq2SeqLM.from_pretrained(self.config.model).to(torch.device(self.config.device))
-            self.tokenizer = AutoTokenizer.from_pretrained(self.config.model, use_fast=self.config.fast_tokenizer)
-            logging.info(f"Model and tokenizer loaded successfully on device '{self.config.device}': {self.config.model}")
+            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
@@ -73,6 +78,13 @@ class SummarizationTaskService(ModelTaskService):
     async def _shutdown(self) -> None:
         self.model = None
         self.tokenizer = None
+        self.device = None
 
     async def _run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
-        return await SummarizationTaskAction(action, self.model, self.tokenizer).run(context)
+        return await SummarizationTaskAction(action, self.model, self.tokenizer, self.device).run(context)
+
+    def _get_model_class(self) -> Type[PreTrainedModel]:
+        return AutoModelForSeq2SeqLM
+
+    def _get_tokenizer_class(self) -> Type[PreTrainedTokenizer]:
+        return AutoTokenizer

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

@@ -7,13 +7,15 @@ 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
 
 class TextClassificationTaskAction:
-    def __init__(self, config: TextClassificationModelActionConfig, model: PreTrainedModel, tokenizer: PreTrainedTokenizer):
+    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)
@@ -26,7 +28,8 @@ class TextClassificationTaskAction:
 
         for index in range(0, len(texts), batch_size):
             batch_texts = texts[index:index + batch_size]
-            inputs = self.tokenizer(batch_texts, return_tensors="pt", padding=True, truncation=True).to(self.model.device)
+            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.no_grad():
                 outputs: SequenceClassifierOutput = self.model(**inputs)
@@ -60,12 +63,14 @@ class TextClassificationTaskService(ModelTaskService):
 
         self.model: Optional[PreTrainedModel] = None
         self.tokenizer: Optional[PreTrainedTokenizer] = None
+        self.device: Optional[torch.device] = None
 
     async def _serve(self) -> None:
         try:
-            self.model = AutoModelForSequenceClassification.from_pretrained(self.config.model).to(torch.device(self.config.device))
-            self.tokenizer = AutoTokenizer.from_pretrained(self.config.model, use_fast=self.config.fast_tokenizer)
-            logging.info(f"Model and tokenizer loaded successfully on device '{self.config.device}': {self.config.model}")
+            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
@@ -73,6 +78,13 @@ class TextClassificationTaskService(ModelTaskService):
     async def _shutdown(self) -> None:
         self.model = None
         self.tokenizer = None
+        self.device = None
 
     async def _run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
-        return await TextClassificationTaskAction(action, self.model, self.tokenizer).run(context, self.config.labels)
+        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.2 → model_compose-0.3.3}/src/mindor/core/component/services/model/tasks/text_embedding.py

@@ -10,10 +10,11 @@ from torch import Tensor
 import torch
 
 class TextEmbeddingTaskAction:
-    def __init__(self, config: TextEmbeddingModelActionConfig, model: PreTrainedModel, tokenizer: PreTrainedTokenizer):
+    def __init__(self, config: TextEmbeddingModelActionConfig, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, device: torch.device):
         self.config: TextEmbeddingModelActionConfig = config
         self.model: PreTrainedModel = model
         self.tokenizer: PreTrainedTokenizer = tokenizer
+        self.device: torch.device = device
 
     async def run(self, context: ComponentActionContext) -> Any:
         text: Union[str, List[str]] = await context.render_variable(self.config.text)
@@ -28,13 +29,14 @@ class TextEmbeddingTaskAction:
 
         for index in range(0, len(texts), batch_size):
             batch_texts = texts[index:index + batch_size]
-            inputs = self.tokenizer(batch_texts, return_tensors="pt", max_length=max_input_length, padding=True, truncation=True).to(self.model.device)
-            attention_mask: Tensor = inputs.get("attention_mask", None)
+            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:
@@ -75,12 +77,14 @@ class TextEmbeddingTaskService(ModelTaskService):
 
         self.model: Optional[PreTrainedModel] = None
         self.tokenizer: Optional[PreTrainedTokenizer] = None
+        self.device: Optional[torch.device] = None
 
     async def _serve(self) -> None:
         try:
-            self.model = AutoModel.from_pretrained(self.config.model).to(torch.device(self.config.device))
-            self.tokenizer = AutoTokenizer.from_pretrained(self.config.model, use_fast=self.config.fast_tokenizer)
-            logging.info(f"Model and tokenizer loaded successfully on device '{self.config.device}': {self.config.model}")
+            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
@@ -88,6 +92,13 @@ class TextEmbeddingTaskService(ModelTaskService):
     async def _shutdown(self) -> None:
         self.model = None
         self.tokenizer = None
+        self.device = None
 
     async def _run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
-        return await TextEmbeddingTaskAction(action, self.model, self.tokenizer).run(context)
+        return await TextEmbeddingTaskAction(action, self.model, self.tokenizer, self.device).run(context)
+
+    def _get_model_class(self) -> Type[PreTrainedModel]:
+        return AutoModel
+
+    def _get_tokenizer_class(self) -> Type[PreTrainedTokenizer]:
+        return AutoTokenizer

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

@@ -5,13 +5,15 @@ 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 torch import Tensor
 import torch
 
 class TextGenerationTaskAction:
-    def __init__(self, config: TextGenerationModelActionConfig, model: Union[PreTrainedModel, GenerationMixin], tokenizer: PreTrainedTokenizer):
+    def __init__(self, config: TextGenerationModelActionConfig, model: Union[PreTrainedModel, GenerationMixin], tokenizer: PreTrainedTokenizer, device: torch.device):
         self.config: TextGenerationModelActionConfig = config
         self.model: Union[PreTrainedModel, GenerationMixin] = model
         self.tokenizer: PreTrainedTokenizer = tokenizer
+        self.device: torch.device = device
 
     async def run(self, context: ComponentActionContext) -> Any:
         prompt: Union[str, List[str]] = await context.render_variable(self.config.prompt)
@@ -28,7 +30,8 @@ class TextGenerationTaskAction:
 
         for index in range(0, len(prompts), batch_size):
             batch_prompts = prompts[index:index + batch_size]
-            inputs = self.tokenizer(batch_prompts, return_tensors="pt").to(self.model.device)
+            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(
@@ -56,12 +59,14 @@ class TextGenerationTaskService(ModelTaskService):
 
         self.model: Optional[Union[PreTrainedModel, GenerationMixin]] = None
         self.tokenizer: Optional[PreTrainedTokenizer] = None
+        self.device: Optional[torch.device] = None
 
     async def _serve(self) -> None:
         try:
-            self.model = AutoModelForCausalLM.from_pretrained(self.config.model).to(torch.device(self.config.device))
-            self.tokenizer = AutoTokenizer.from_pretrained(self.config.model, use_fast=self.config.fast_tokenizer)
-            logging.info(f"Model and tokenizer loaded successfully on device '{self.config.device}': {self.config.model}")
+            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
@@ -69,6 +74,13 @@ class TextGenerationTaskService(ModelTaskService):
     async def _shutdown(self) -> None:
         self.model = None
         self.tokenizer = None
+        self.device = None
 
     async def _run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
-        return await TextGenerationTaskAction(action, self.model, self.tokenizer).run(context)
+        return await TextGenerationTaskAction(action, self.model, self.tokenizer, self.device).run(context)
+
+    def _get_model_class(self) -> Type[PreTrainedModel]:
+        return AutoModelForCausalLM
+
+    def _get_tokenizer_class(self) -> Type[PreTrainedTokenizer]:
+        return AutoTokenizer

{model_compose-0.3.2 → model_compose-0.3.3}/src/mindor/core/component/services/model/tasks/translation.py

@@ -5,13 +5,15 @@ 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 torch import Tensor
 import torch
 
 class TranslationTaskAction:
-    def __init__(self, config: TranslationModelActionConfig, model: Union[PreTrainedModel, GenerationMixin], tokenizer: PreTrainedTokenizer):
+    def __init__(self, config: TranslationModelActionConfig, model: Union[PreTrainedModel, GenerationMixin], tokenizer: PreTrainedTokenizer, device: torch.device):
         self.config: TranslationModelActionConfig = config
         self.model: Union[PreTrainedModel, GenerationMixin] = model
         self.tokenizer: PreTrainedTokenizer = tokenizer
+        self.device: torch.device = device
 
     async def run(self, context: ComponentActionContext) -> Any:
         text: Union[str, List[str]] = await context.render_variable(self.config.text)
@@ -28,8 +30,9 @@ class TranslationTaskAction:
 
         for index in range(0, len(texts), batch_size):
             batch_texts = texts[index:index + batch_size]
-            inputs = self.tokenizer(batch_texts, return_tensors="pt", max_length=max_input_length, padding=True, truncation=True).to(self.model.device)
-
+            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,
@@ -57,12 +60,14 @@ class TranslationTaskService(ModelTaskService):
 
         self.model: Optional[Union[PreTrainedModel, GenerationMixin]] = None
         self.tokenizer: Optional[PreTrainedTokenizer] = None
+        self.device: Optional[torch.device] = None
 
     async def _serve(self) -> None:
         try:
-            self.model = AutoModelForSeq2SeqLM.from_pretrained(self.config.model).to(torch.device(self.config.device))
-            self.tokenizer = AutoTokenizer.from_pretrained(self.config.model, use_fast=self.config.fast_tokenizer)
-            logging.info(f"Model and tokenizer loaded successfully on device '{self.config.device}': {self.config.model}")
+            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
@@ -70,6 +75,13 @@ class TranslationTaskService(ModelTaskService):
     async def _shutdown(self) -> None:
         self.model = None
         self.tokenizer = None
+        self.device = None
 
     async def _run(self, action: ModelActionConfig, context: ComponentActionContext) -> Any:
-        return await TranslationTaskAction(action, self.model, self.tokenizer).run(context)
+        return await TranslationTaskAction(action, self.model, self.tokenizer, self.device).run(context)
+
+    def _get_model_class(self) -> Type[PreTrainedModel]:
+        return AutoModelForSeq2SeqLM
+
+    def _get_tokenizer_class(self) -> Type[PreTrainedTokenizer]:
+        return AutoTokenizer

{model_compose-0.3.2 → model_compose-0.3.3}/src/mindor/core/controller/base.py

@@ -17,7 +17,7 @@ from mindor.core.workflow import Workflow, WorkflowResolver, create_workflow
 from mindor.core.logger import LoggerService, create_logger
 from mindor.core.controller.webui import ControllerWebUI
 from mindor.core.utils.workqueue import WorkQueue
-from mindor.core.utils.expiring import ExpiringDict
+from mindor.core.utils.caching import ExpiringDict
 from .runtime.specs import ControllerRuntimeSpecs
 from .runtime.docker import DockerRuntimeLauncher
 from threading import Lock