Fixed import to ensure that tests run

Author: KennethEnevoldsen

.gitignore

@@ -162,3 +162,6 @@ test.py
 tmp.ipynb
 tmp.py
 *.pickle
+
+# vscode
+.vscode/

tests/test_chatgpt.py

@@ -1,15 +1,11 @@
 import unittest
 from unittest.mock import patch
 
-from skllm.openai.chatgpt import (
-    construct_message,
-    extract_json_key,
-    get_chat_completion,
-)
+from skllm.openai.chatgpt import construct_message, get_chat_completion
+from skllm.utils import extract_json_key
 
 
 class TestChatGPT(unittest.TestCase):
-
     @patch("skllm.openai.credentials.set_credentials")
     @patch("openai.ChatCompletion.create")
     def test_get_chat_completion(self, mock_create, mock_set_credentials):
@@ -21,9 +17,18 @@ def test_get_chat_completion(self, mock_create, mock_set_credentials):
 
         result = get_chat_completion(messages, key, org, model)
 
-        self.assertTrue(mock_set_credentials.call_count <= 1, "set_credentials should be called at most once")
-        self.assertEqual(mock_create.call_count, 2, "ChatCompletion.create should be called twice due to an exception "
-                                                    "on the first call")
+        self.assertTrue(
+            mock_set_credentials.call_count <= 1,
+            "set_credentials should be called at most once",
+        )
+        self.assertEqual(
+            mock_create.call_count,
+            2,
+            (
+                "ChatCompletion.create should be called twice due to an exception "
+                "on the first call"
+            ),
+        )
         self.assertEqual(result, "success")
 
     def test_construct_message(self):
@@ -45,5 +50,5 @@ def test_extract_json_key(self):
         self.assertEqual(result_with_invalid_key, None)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

tests/test_gpt_few_shot_clf.py

@@ -5,8 +5,7 @@
 
 import numpy as np
 
-from skllm.models.gpt_dyn_few_shot_clf import DynamicFewShotGPTClassifier
-from skllm.models.gpt_few_shot_clf import FewShotGPTClassifier
+from skllm import DynamicFewShotGPTClassifier, FewShotGPTClassifier
 
 
 class TestFewShotGPTClassifier(unittest.TestCase):

tests/test_gpt_zero_shot_clf.py

@@ -5,10 +5,7 @@
 
 import numpy as np
 
-from skllm.models.gpt_zero_shot_clf import (
-    MultiLabelZeroShotGPTClassifier,
-    ZeroShotGPTClassifier,
-)
+from skllm import MultiLabelZeroShotGPTClassifier, ZeroShotGPTClassifier
 
 
 def _get_ret(label):