Merge pull request #24 from aixplain/ENG-478-Fixing-model-interfaces-…

…inconsistencies

Correcting abstract method inconsistencies

CHANGELOG.md

@@ -3,8 +3,17 @@ All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](http://keepachangelog.com/)
 and this project adheres to [Semantic Versioning](http://semver.org/).
+
+## [1.0.0] - 2024-08-28
 
-## [v0.0.2] - 2024-07-25
+### Added
+- Test for TextGenerationChatModel.
+### Changed
+
+### Fixed
+- Added @abstractmethod decorators to all abstract functions in function_models and aixplain_models.
+
+## [0.0.2] - 2024-07-25
 
 ### Added
 - Added support for script nodes.

aixplain/model_interfaces/__version__.py

@@ -1,7 +1,7 @@
 __title__ = "model-interfaces"
 __description__ = "model-interfaces is the interface to host your models on aiXplain"
 __url__ = "https://github.com/aixplain/aixplain-models/tree/main/docs"
-__version__ = "0.0.2"
+__version__ = "1.0.0"
 __author__ = "Duraikrishna Selvaraju and Michael Lam"
 __author_email__ = "[email protected]"
 __license__ = "http://www.apache.org/licenses/LICENSE-2.0"

aixplain/model_interfaces/interfaces/aixplain_model.py

@@ -1,12 +1,16 @@
 from kserve.model import Model
 from typing import Dict, List
+from abc import abstractmethod
+
 from aixplain.model_interfaces.schemas.function.function_input import APIInput
 from aixplain.model_interfaces.schemas.function.function_output import APIOutput
 
 class AixplainModel(Model):
 
+    @abstractmethod
     def run_model(self, api_input: Dict[str, List[APIInput]], headers: Dict[str, str] = None) -> Dict[str, List[APIOutput]]:
         pass
 
+    @abstractmethod
     def predict(self, request: Dict[str, List[APIInput]], headers: Dict[str, str] = None) -> Dict[str, List[APIOutput]]:
         pass

pyproject.toml

@@ -15,7 +15,7 @@ includes = ["aixplain"]
 
 [project]
 name = "model-interfaces"
-version = "0.0.2"
+version = "1.0.0"
 description = "A package specifying the model interfaces supported by aiXplain"
 license = { text = "Apache License, Version 2.0: http://www.apache.org/licenses/LICENSE-2.0" }
 dependencies = [

tests/unit_tests/models/test_mock_classification.py

@@ -26,13 +26,13 @@ def test_predict(self):
         predict_output = mock_model.predict(predict_input)
         output_dict = predict_output["predictions"][0]
 
-        assert output_dict["data"] == "positive"
-        assert output_dict["predicted_labels"][0]["label"] == "positive"
-        assert output_dict["predicted_labels"][0]["confidence"] == 0.7
+        assert output_dict.data == "positive"
+        assert output_dict.predicted_labels[0].label == "positive"
+        assert output_dict.predicted_labels[0].confidence == 0.7
 
 class MockModel(ClassificationModel):
-    def run_model(self, api_input: Dict[str, List[ClassificationInput]], headers: Dict[str, str] = None) -> Dict[str, List[ClassificationOutput]]:
-        instances = api_input["instances"]
+    def run_model(self, api_input: List[ClassificationInput], headers: Dict[str, str] = None) -> List[ClassificationOutput]:
+        instances = api_input
         predictions_list = []
         # There's only 1 instance in this case.
         for instance in instances:
@@ -50,7 +50,6 @@ def run_model(self, api_input: Dict[str, List[ClassificationInput]], headers: Di
                 "data": data,
                 "predicted_labels": labels
             }
-            speech_recognition_output = ClassificationOutput(**output_dict)
-            predictions_list.append(speech_recognition_output)
-        predict_output = {"predictions": predictions_list}
-        return predict_output
+            classification_output = ClassificationOutput(**output_dict)
+            predictions_list.append(classification_output)
+        return predictions_list

tests/unit_tests/models/test_mock_diacritization.py

@@ -26,16 +26,15 @@ def test_predict(self):
         predict_output = mock_model.predict(predict_input)
         output_dict = predict_output["predictions"][0]
 
-        assert output_dict["data"] == "السَّلَامُ عَلَيْكُمْ"
-        assert output_dict["details"]["text"] == "السَّلَامُ عَلَيْكُمْ"
-        assert output_dict["details"]["confidence"] == 0.7
+        assert output_dict.data == "السَّلَامُ عَلَيْكُمْ"
+        assert output_dict.details.text == "السَّلَامُ عَلَيْكُمْ"
+        assert output_dict.details.confidence == 0.7
 
 class MockModel(DiacritizationModel):
-    def run_model(self, api_input: Dict[str, List[DiacritizationInput]], headers: Dict[str, str] = None) -> Dict[str, List[DiacritizationOutput]]:
-        instances = api_input["instances"]
+    def run_model(self, api_input: List[DiacritizationInput], headers: Dict[str, str] = None) -> List[DiacritizationOutput]:
         predictions_list = []
         # There's only 1 instance in this case.
-        for instance in instances:
+        for instance in api_input:
             instance_data = instance.dict()
             model_instance = Mock()
             model_instance.process_data.return_value = ("السَّلَامُ عَلَيْكُمْ", 0.7)
@@ -52,5 +51,4 @@ def run_model(self, api_input: Dict[str, List[DiacritizationInput]], headers: Di
             }
             speech_recognition_output = DiacritizationOutput(**output_dict)
             predictions_list.append(speech_recognition_output)
-        predict_output = {"predictions": predictions_list}
-        return predict_output
+        return predictions_list

tests/unit_tests/models/test_mock_fill_text_mask.py

@@ -31,14 +31,13 @@ def test_predict(self):
         predict_output = mock_model.predict(predict_input)
         output_dict = predict_output["predictions"][0]
 
-        assert output_dict["data"] == "We are filling a text mask."
+        assert output_dict.data == "We are filling a text mask."
 
 class MockModel(FillTextMaskModel):
-    def run_model(self, api_input: Dict[str, List[FillTextMaskInput]], headers: Dict[str, str] = None) -> Dict[str, List[FillTextMaskOutput]]:
-        instances = api_input["instances"]
+    def run_model(self, api_input: List[FillTextMaskInput], headers: Dict[str, str] = None) -> List[FillTextMaskOutput]:
         predictions_list = []
         # There's only 1 instance in this case.
-        for instance in instances:
+        for instance in api_input:
             instance_data = instance.dict()
             model_instance = Mock()
             model_instance.process_data.return_value = "We are filling a text mask."
@@ -53,5 +52,4 @@ def run_model(self, api_input: Dict[str, List[FillTextMaskInput]], headers: Dict
             }
             output = FillTextMaskOutput(**output_dict)
             predictions_list.append(output)
-        predict_output = {"predictions": predictions_list}
-        return predict_output
+        return predictions_list

tests/unit_tests/models/test_mock_speech_enhancement.py

@@ -35,14 +35,13 @@ def test_predict(self):
         predict_output = mock_model.predict(predict_input)
         output_dict = predict_output["predictions"][0]
 
-        assert output_dict["data"] == "VGhpcyBpcyBhbiBhdWRpbyBvdXRwdXQ="
+        assert output_dict.data == "VGhpcyBpcyBhbiBhdWRpbyBvdXRwdXQ="
 
 class MockModel(SpeechEnhancementModel):
-    def run_model(self, api_input: Dict[str, List[SpeechEnhancementInput]], headers: Dict[str, str] = None) -> Dict[str, List[SpeechEnhancementOutput]]:
-        instances = api_input["instances"]
+    def run_model(self, api_input: List[SpeechEnhancementInput], headers: Dict[str, str] = None) -> List[SpeechEnhancementOutput]:
         predictions_list = []
         # There's only 1 instance in this case.
-        for instance in instances:
+        for instance in api_input:
             instance_data = instance.dict()
             model_instance = Mock()
             model_instance.process_data.return_value = encode(b"This is an audio output")
@@ -59,5 +58,4 @@ def run_model(self, api_input: Dict[str, List[SpeechEnhancementInput]], headers:
             }
             speech_recognition_output = SpeechEnhancementOutput(**output_dict)
             predictions_list.append(speech_recognition_output)
-        predict_output = {"predictions": predictions_list}
-        return predict_output
+        return predictions_list

tests/unit_tests/models/test_mock_speech_recognition.py

@@ -35,16 +35,15 @@ def test_predict(self):
         predict_output = mock_model.predict(predict_input)
         output_dict = predict_output["predictions"][0]
 
-        assert output_dict["data"] == "This is a test transcription"
-        assert output_dict["details"]["text"] == "This is a test transcription"
-        assert output_dict["details"]["confidence"] == 0.7
+        assert output_dict.data == "This is a test transcription"
+        assert output_dict.details.text == "This is a test transcription"
+        assert output_dict.details.confidence == 0.7
 
 class MockModel(SpeechRecognitionModel):
-    def run_model(self, api_input: Dict[str, List[SpeechRecognitionInput]], headers: Dict[str, str] = None) -> Dict[str, List[SpeechRecognitionOutput]]:
-        instances = api_input["instances"]
+    def run_model(self, api_input: List[SpeechRecognitionInput], headers: Dict[str, str] = None) -> List[SpeechRecognitionOutput]:
         predictions_list = []
         # There's only 1 instance in this case.
-        for instance in instances:
+        for instance in api_input:
             instance_data = instance.dict()
             model_instance = Mock()
             model_instance.process_data.return_value = ("This is a test transcription", 0.7)
@@ -61,5 +60,4 @@ def run_model(self, api_input: Dict[str, List[SpeechRecognitionInput]], headers:
             }
             speech_recognition_output = SpeechRecognitionOutput(**output_dict)
             predictions_list.append(speech_recognition_output)
-        predict_output = {"predictions": predictions_list}
-        return predict_output
+        return predictions_list

tests/unit_tests/models/test_mock_subtitle_translation.py

@@ -33,14 +33,13 @@ def test_predict(self):
         predict_output = mock_model.predict(predict_input)
         output_dict = predict_output["predictions"][0]
 
-        assert output_dict["data"] == "This is a subtitle translation."
+        assert output_dict.data == "This is a subtitle translation."
 
 class MockModel(SubtitleTranslationModel):
-    def run_model(self, api_input: Dict[str, List[SubtitleTranslationInput]], headers: Dict[str, str] = None) -> Dict[str, List[SubtitleTranslationOutput]]:
-        instances = api_input["instances"]
+    def run_model(self, api_input: List[SubtitleTranslationInput], headers: Dict[str, str] = None) -> List[SubtitleTranslationOutput]:
         predictions_list = []
         # There's only 1 instance in this case.
-        for instance in instances:
+        for instance in api_input:
             instance_data = instance.dict()
             model_instance = Mock()
             model_instance.process_data.return_value = "This is a subtitle translation."
@@ -55,5 +54,4 @@ def run_model(self, api_input: Dict[str, List[SubtitleTranslationInput]], header
             }
             search_output = SubtitleTranslationOutput(**output_dict)
             predictions_list.append(search_output)
-        predict_output = {"predictions": predictions_list}
-        return predict_output
+        return predictions_list

tests/unit_tests/models/test_mock_text_generation.py

@@ -2,6 +2,7 @@
 from aixplain.model_interfaces.schemas.function.function_input import TextGenerationInput
 from aixplain.model_interfaces.schemas.function.function_output import TextGenerationOutput
 from aixplain.model_interfaces.interfaces.function_models import TextGenerationModel
+from aixplain.model_interfaces.schemas.modality.modality_input import TextListInput
 from typing import Dict, List
 
 class TestMockTextGeneration():
@@ -25,6 +26,21 @@ def test_predict(self):
 
         assert predictions.data == "I am a text generation model."
 
+    def test_tokenize(self):
+        tokenize_input = {
+            # provide a list of test instances
+            "instances": [
+                {
+                    "data": ["Hello world", "Hello world again"]
+                }
+            ],
+            "function": "TOKENIZE"
+        }
+        mock_model = MockModel("Mock")
+        token_counts_list = mock_model.predict(tokenize_input)
+
+        assert token_counts_list["token_counts"][0] == [11, 17]
+
 class MockModel(TextGenerationModel):
     def run_model(self, api_input: List[TextGenerationInput], headers: Dict[str, str] = None) -> List[TextGenerationOutput]:
         print(f"API INPUT: {api_input}")
@@ -47,4 +63,11 @@ def run_model(self, api_input: List[TextGenerationInput], headers: Dict[str, str
             text_generation_output = TextGenerationOutput(**output_dict)
             predictions_list.append(text_generation_output)
         predict_output = predictions_list
-        return predict_output
+        return predict_output
+
+    def tokenize(self, api_input: List[TextListInput], headers: Dict[str, str] = None) -> List[List[int]]:
+        token_counts_list = []
+        for instance in api_input:
+            token_counts = [len(message) for message in instance.data]
+            token_counts_list.append(token_counts)
+        return token_counts_list

tests/unit_tests/models/test_mock_text_generation_chat.py

@@ -0,0 +1,118 @@
+from unittest.mock import Mock
+from aixplain.model_interfaces.schemas.function.function_input import TextGenerationInput
+from aixplain.model_interfaces.schemas.function.function_output import TextGenerationOutput
+from aixplain.model_interfaces.interfaces.function_models import TextGenerationChatModel, TextGenerationChatTemplatizeInput
+from aixplain.model_interfaces.schemas.modality.modality_input import TextListInput
+from typing import Dict, List, Text
+
+class TestMockTextGeneration():
+    def test_predict(self):
+        predict_input = {
+            "instances": [
+                {
+                    "data": "How many cups in a liter?",
+                    "max_new_tokens": 200,
+                    "top_p": 0.92,
+                    "top_k": 1,
+                    "num_return_sequences": 1
+                }
+            ],
+            "function": "predict"
+        }
+
+        mock_model = MockModel("Mock")
+        predict_output = mock_model.predict(predict_input)
+        predictions = predict_output["predictions"][0]
+
+        assert predictions.data == "I am a text generation model."
+
+    def test_tokenize(self):
+        tokenize_input = {
+            # provide a list of test instances
+            "instances": [
+                {
+                    "data": ["Hello world", "Hello world again"]
+                }
+            ],
+            "function": "TOKENIZE"
+        }
+        mock_model = MockModel("Mock")
+        token_counts_list = mock_model.predict(tokenize_input)
+        print(f"Token counts: {token_counts_list}")
+
+        assert token_counts_list["token_counts"][0] == [11, 17]
+
+    def test_templatize(self):
+        data_to_be_templatized = [
+            {
+                "role": "user",
+                "content": "Hello, how are you?"
+            },
+            {
+                "role": "assistant",
+                "content": "I'm doing great. How can I help you today?"
+            },
+            {
+                "role": "user",
+                "content": "I'd like to show off how chat templating works!"
+            },
+            {
+                "role": "system",
+                "content": "I'd like to show off how chat templating works!"
+            }
+        ]
+        templatize_input = {
+            "instances": [
+                {
+                    "data": data_to_be_templatized
+                }
+            ],
+            "function": "TEMPLATIZE"
+        }
+
+        mock_model = MockModel("Mock")
+        templatized_text = mock_model.predict(templatize_input)
+
+        assert templatized_text["prompts"][0] == f"Mock template: {str(data_to_be_templatized)}"
+        # for i in range(len(data_to_be_templatized)):
+        #     print(f"templatized_text: {templatized_text}")
+        #     assert templatized_text["prompts"][i] == f"Mock template: {str(data_to_be_templatized[i])}"
+
+
+class MockModel(TextGenerationChatModel):
+    def run_model(self, api_input: List[TextGenerationInput], headers: Dict[str, str] = None) -> List[TextGenerationOutput]:
+        print(f"API INPUT: {api_input}")
+        instances = api_input
+        predictions_list = []
+        # There's only 1 instance in this case.
+        for instance in instances:
+            instance_data = instance.dict()
+            model_instance = Mock()
+            model_instance.process_data.return_value = "I am a text generation model."
+            result = model_instance.process_data(instance_data["data"])
+            model_instance.delete()
+
+            # Map back onto TextGenerationOutput
+            data = result
+
+            output_dict = {
+                "data": data,
+            }
+            text_generation_output = TextGenerationOutput(**output_dict)
+            predictions_list.append(text_generation_output)
+        predict_output = predictions_list
+        return predict_output
+
+    def tokenize(self, api_input: List[TextListInput], headers: Dict[str, str] = None) -> List[List[int]]:
+        token_counts_list = []
+        for instance in api_input:
+            token_counts = [len(message) for message in instance.data]
+            token_counts_list.append(token_counts)
+        return token_counts_list
+
+    def templatize(self, api_input: List[TextGenerationChatTemplatizeInput], headers: Dict[str, str] = None) -> List[Text]:
+        template_text_list = []
+        for instance in api_input:
+            templatized_text = f"Mock template: {str(instance.data)}"
+            template_text_list.append(templatized_text)
+        return template_text_list