✨ introduce request dm fetching in runtime

Signed-off-by: Prashant Gupta <[email protected]>

caikit/runtime/__init__.py

@@ -11,3 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+# Local
+from .service_factory import get_request, get_train_params, get_train_request

caikit/runtime/service_factory.py

@@ -31,6 +31,7 @@
 # Local
 from caikit import get_config
 from caikit.core import LocalBackend, ModuleBase, registries
+from caikit.core.data_model.base import DataBase
 from caikit.core.data_model.dataobject import _AUTO_GEN_PROTO_CLASSES
 from caikit.interfaces.runtime.data_model import (
     TrainingInfoRequest,
@@ -266,3 +267,38 @@ def _get_and_filter_modules(
             excluded_modules,
         )
         return clean_modules
+
+
+def get_request(
+    module_class: Type[ModuleBase],
+    input_streaming: bool = False,
+    output_streaming: bool = False,
+) -> Type[DataBase]:
+    """Helper function to return the request DataModel for the Module Class"""
+    if input_streaming and output_streaming:
+        request_class_name = f"BidiStreaming{module_class.TASK_CLASS.__name__}Request"
+    elif input_streaming:
+        request_class_name = f"ClientStreaming{module_class.TASK_CLASS.__name__}Request"
+    elif output_streaming:
+        request_class_name = f"ServerStreaming{module_class.TASK_CLASS.__name__}Request"
+    else:
+        request_class_name = f"{module_class.TASK_CLASS.__name__}Request"
+    return DataBase.get_class_for_name(request_class_name)
+
+
+def get_train_request(module_class: Type[ModuleBase]) -> Type[DataBase]:
+    """Helper function to return the train request DataModel for the Module Class"""
+    request_class_name = (
+        f"{module_class.TASK_CLASS.__name__}{module_class.__name__}TrainRequest"
+    )
+    print(request_class_name)
+    return DataBase.get_class_for_name(request_class_name)
+
+
+def get_train_params(module_class: Type[ModuleBase]) -> Type[DataBase]:
+    """Helper function to return the train parameters DataModel for the Module Class"""
+    request_class_name = (
+        f"{module_class.TASK_CLASS.__name__}{module_class.__name__}TrainParameters"
+    )
+    print(request_class_name)
+    return DataBase.get_class_for_name(request_class_name)

examples/text-sentiment/client.py

@@ -21,6 +21,7 @@
 
 # Local
 from caikit.config.config import get_config
+from caikit.runtime import get_request
 from caikit.runtime.service_factory import ServicePackageFactory
 from text_sentiment.runtime_model import HuggingFaceSentimentModule
 import caikit
@@ -51,7 +52,7 @@
 
         # Run inference for two sample prompts
         for text in ["I am not feeling well today!", "Today is a nice sunny day"]:
-            request = HuggingFaceSentimentModule.TASK_CLASS.UNARY_REQUEST_DATA_MODEL(
+            request = get_request(HuggingFaceSentimentModule)(
                 text_input=text
             ).to_proto()
             response = client_stub.HuggingFaceSentimentTaskPredict(

tests/fixtures/sample_lib/__init__.py

@@ -3,7 +3,14 @@
 
 # Local
 from . import data_model, modules
-from .modules import InnerModule, OtherModule, SampleModule, SamplePrimitiveModule
+from .modules import (
+    CompositeModule,
+    InnerModule,
+    OtherModule,
+    SampleModule,
+    SamplePrimitiveModule,
+    StreamingModule,
+)
 from caikit.config import configure
 
 # Run configure for sample_lib configuration

tests/fixtures/sample_lib/modules/__init__.py

@@ -7,4 +7,5 @@
     InnerModule,
     SampleModule,
     SamplePrimitiveModule,
+    StreamingModule,
 )