add ut

Signed-off-by: Lv, Liang1 <[email protected]>

Author: lvliang-intel

neural_chat/chatbot.py

@@ -21,7 +21,7 @@
 from .config import FinetuningConfig
 from .pipeline.finetuning.finetuning import Finetuning
 from .config import DeviceOptions, BackendOptions, AudioOptions
-from models.base_model import get_model_adapter
+from .models.base_model import get_model_adapter
 from .utils.common import get_device_type, get_backend_type
 from .pipeline.plugins.caching.cache import init_similar_cache_from_config
 from .pipeline.plugins.audio.asr import AudioSpeechRecognition
@@ -35,8 +35,11 @@ def build_chatbot(config: NeuralChatConfig):
     """Build the chatbot with a given configuration.
 
     Args:
-        config (NeuralChatConfig):  The class of NeuralChatConfig containing model path,
-                                    device, backend, plugin config etc.
+        config (NeuralChatConfig): Configuration for building the chatbot.
+
+    Returns:
+        adapter: The chatbot model adapter.
+
     Example:
         from neural_chat.config import NeuralChatConfig
         from neural_chat.chatbot import build_chatbot
@@ -45,7 +48,7 @@ def build_chatbot(config: NeuralChatConfig):
         response = chatbot.predict("Tell me about Intel Xeon Scalable Processors.")
     """
     # Validate input parameters
-    if config.device not in [option.name for option in DeviceOptions]:
+    if config.device not in [option.name.lower() for option in DeviceOptions]:
         valid_options = ", ".join([option.name.lower() for option in DeviceOptions])
         raise ValueError(f"Invalid device value '{config.device}'. Must be one of {valid_options}")
 
@@ -111,10 +114,20 @@ def build_chatbot(config: NeuralChatConfig):
     return adapter
 
 def finetune_model(config: FinetuningConfig):
+    """Finetune the model based on the provided configuration.
+
+    Args:
+        config (FinetuningConfig): Configuration for finetuning the model.
+    """
+
     assert config is not None, "FinetuningConfig is needed for finetuning."
     finetuning = Finetuning(config)
     finetuning.finetune()
 
 def optimize_model(config: OptimizationConfig):
-    # Implement the logic to optimize the model
+    """Optimize the model based on the provided configuration.
+
+    Args:
+        config (OptimizationConfig): Configuration for optimizing the model.
+    """
     pass

neural_chat/requirements.txt

@@ -1,5 +1,5 @@
 transformers>=4.31.0
-perft
+peft
 fastchat
 torch
 intel_extension_for_pytorch
@@ -8,3 +8,5 @@ speechbrain
 paddlepaddle
 paddlespeech
 shortuuid
+gptcache
+librosa>=0.10.0

neural_chat/tests/api/test_chatbot_build_api.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright (c) 2023 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+import unittest
+from neural_chat.chatbot import build_chatbot
+from neural_chat.config import NeuralChatConfig
+
+class TestChatbotBuilder(unittest.TestCase):
+    def setUp(self):
+        pass
+
+    def test_build_chatbot_valid_config(self):
+        config = NeuralChatConfig()
+        config.device = "cpu"
+        config.backend = "torch"
+        config.model_name_or_path = "meta-llama/Llama-2-7b-chat-hf"
+
+        chatbot = build_chatbot(config)
+        self.assertIsNotNone(chatbot)
+        response = chatbot.predict("Tell me about Intel Xeon Scalable Processors.")
+        print(response)
+
+    def test_build_chatbot_invalid_config(self):
+        # Similar to the previous test, but with an invalid configuration
+        pass
+
+    def test_build_chatbot_retrieval(self):
+        # Test the retrieval logic
+        pass
+
+    def test_build_chatbot_audio(self):
+        # Test the audio logic
+        pass
+
+    # Add more tests for other components of the function
+
+if __name__ == '__main__':
+    unittest.main()

neural_chat/utils/common.py

@@ -16,12 +16,18 @@
 # limitations under the License.
 
 import torch
-import habana_frameworks.torch.hpu as hthpu
+
+try:
+    import habana_frameworks.torch.hpu as hthpu
+    is_hpu_available = True
+except ImportError:
+    print("Package 'habana_frameworks.torch.hpu' is not installed.")
+    is_hpu_available = False
 
 def get_device_type():
     if torch.cuda.is_available():
         device = "cuda"
-    elif hthpu.is_available():
+    elif is_hpu_available:
         device = "hpu"
     elif torch.xpu.is_available():
         device = "xpu"