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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: model-compose
-Version: 0.3.2
+Version: 0.3.4
 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.4}/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.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "model-compose"
-version = "0.3.2"
+version = "0.3.4"
 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.4}/src/mindor/cli/compose.py

@@ -11,11 +11,11 @@ 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:
-    """model-compose (from Mindor)"""
+    """model-compose"""
     ctx.ensure_object(dict)
     ctx.obj["config_files"] = list(config_files)
 
@@ -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
@@ -138,7 +138,7 @@ def stop_command(
     asyncio.run(_async_command())
 
 @click.command(name="run")
-@click.argument("workflow", required=False)
+@click.argument("workflow_id", required=False)
 @click.option(
     "--input", "-i", "input_json",
     type=str,
@@ -149,20 +149,27 @@ 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(
+    "--output", "-o", "output_path",
+    type=click.Path(dir_okay=False, writable=True, path_type=Path),
+    required=False,
+    help="Path to save the output result as a file."
 )
 @click.option("--verbose", "-v", is_flag=True, help="Enable verbose output.")
 @click.pass_context
 def run_command(
     ctx: click.Context,
-    workflow: Optional[str],
+    workflow_id: Optional[str],
     input_json: Optional[str],
     env_files: List[Path],
     env_data: List[str],
+    output_path: Optional[Path],
     verbose: bool
 ) -> None:
     config_files = ctx.obj.get("config_files", [])
@@ -171,12 +178,16 @@ def run_command(
             env = load_env_files(".", env_files or [])
             config = load_compose_config(".", config_files, env)
             input = json.loads(input_json) if input_json else {}
-            state = await run_workflow(config, workflow, input, verbose)
-            click.echo(json.dumps(
-                state.error if state.error else state.output, 
-                indent=2, 
-                ensure_ascii=False
-            ))
+            state = await run_workflow(config, workflow_id, input, output_path, verbose)
+            if isinstance(state.output, (dict, list)) or state.error:
+                click.echo(json.dumps(
+                    state.output or state.error,
+                    indent=2,
+                    ensure_ascii=False
+                ))
+            else:
+                if state.output is not None:
+                    click.echo(state.output)
         except json.JSONDecodeError:
             click.echo("❌ Invalid JSON provided for --input", err=True)
         except Exception as e:

{model_compose-0.3.2 → model_compose-0.3.4}/src/mindor/core/component/base.py

@@ -78,6 +78,12 @@ class ComponentService(AsyncService):
 
         await super()._stop()
 
+    async def _serve(self) -> None:
+        pass
+
+    async def _shutdown(self) -> None:
+        pass
+
     @abstractmethod
     async def _run(self, action: ActionConfig, context: ComponentActionContext) -> Any:
         pass

{model_compose-0.3.2 → model_compose-0.3.4}/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.2 → model_compose-0.3.4}/src/mindor/core/component/services/__init__.py

@@ -5,3 +5,4 @@ from .mcp_client import *
 from .model import *
 from .workflow import *
 from .shell import *
+from .text_splitter import *

{model_compose-0.3.2 → model_compose-0.3.4}/src/mindor/core/component/services/http_client.py

@@ -113,12 +113,14 @@ class HttpClientComponent(ComponentService):
     def __init__(self, id: str, config: HttpClientComponentConfig, global_configs: ComponentGlobalConfigs, daemon: bool):
         super().__init__(id, config, global_configs, daemon)
 
-        self.client: HttpClient = None
+        self.client: Optional[HttpClient] = None
 
-    async def _serve(self) -> None:
+    async def _start(self) -> None:
         self.client = HttpClient(self.config.base_url, self.config.headers)
+        await super()._start()
 
-    async def _shutdown(self) -> None:
+    async def _stop(self) -> None:
+        await super()._stop()
         await self.client.close()
         self.client = None
 

{model_compose-0.3.2 → model_compose-0.3.4}/src/mindor/core/component/services/mcp_client.py

@@ -32,10 +32,14 @@ class McpClientComponent(ComponentService):
     def __init__(self, id: str, config: McpClientComponentConfig, global_configs: ComponentGlobalConfigs, daemon: bool):
         super().__init__(id, config, global_configs, daemon)
 
-    async def _serve(self) -> None:
+        self.client: Optional[McpClient] = None 
+
+    async def _start(self) -> None:
         self.client = McpClient(self.config.url, self.config.headers)
+        await super()._start()
 
-    async def _shutdown(self) -> None:
+    async def _stop(self) -> None:
+        await super()._stop()
         await self.client.close()
         self.client = None
 

model_compose-0.3.4/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.4}/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.4}/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.4}/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.4}/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.4}/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