fix (#119)

Author: tastelikefeet

assets/slogan.png

@@ -1,11 +1,11 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-from typing import TYPE_CHECKING, List, Union
+from typing import TYPE_CHECKING, Any, List, Union
 
 if TYPE_CHECKING:
     import torch
 
 
 class Patch:
 
-    def __call__(self, module: Union['torch.nn.Module', List['torch.nn.Module']], *args, **kwargs):
+    def __call__(self, module: Union['torch.nn.Module', List['torch.nn.Module'], Any], *args, **kwargs):
         ...

src/twinkle/patch/base.py

@@ -6,6 +6,7 @@
 import twinkle
 from twinkle import remote_function
 from twinkle.data_format import InputFeature, SampleResponse, SamplingParams, Trajectory
+from twinkle.patch import Patch
 from twinkle.template import Template
 from twinkle.utils import construct_class
 
@@ -42,6 +43,10 @@ def sample(
         """
         pass
 
+    @abstractmethod
+    def apply_patch(self, patch_cls: Union[Patch, Type[Patch], str], **kwargs) -> None:
+        ...
+
     @staticmethod
     def _not_encoded(inputs: Any) -> bool:
         """Check if inputs are not yet encoded (i.e., is Trajectory, not InputFeature).

src/twinkle/sampler/base.py

@@ -24,11 +24,12 @@
 import numpy as np
 import os
 import threading
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
 from twinkle import DeviceMesh, get_logger, remote_class, remote_function, requires
 from twinkle.checkpoint_engine import CheckpointEngineMixin
 from twinkle.data_format import InputFeature, SampledSequence, SampleResponse, SamplingParams, Trajectory
+from twinkle.patch import Patch, apply_patch
 from twinkle.patch.vllm_lora_weights import VLLMLoraWeights
 from twinkle.sampler.base import Sampler
 from twinkle.utils import Platform
@@ -212,6 +213,9 @@ def encode_trajectory_for_vllm(self,
             result['videos'] = videos
         return result
 
+    def apply_patch(self, patch_cls: Union[Patch, Type[Patch], str], **kwargs) -> None:
+        apply_patch(self, patch_cls, **kwargs)
+
     async def _sample_single(
         self,
         feat: Dict[str, Any],