switch openai to lazy import (#131)

Author: liusong1222

angelslim/compressor/speculative/train/data/data_generation.py

@@ -7,7 +7,8 @@
 
 import tqdm
 from datasets import load_dataset
-from openai import OpenAI
+
+from angelslim.utils.lazy_imports import openai
 
 from .data_utils import convert_sharegpt_data, convert_ultrachat_data
 
@@ -36,7 +37,7 @@ def _initialize_clients(self, base_port: int, max_clients: int) -> None:
         """Initialize available OpenAI clients."""
         for i in range(max_clients):
             base_url = f"http://localhost:{base_port + i}/v1"
-            client = OpenAI(base_url=base_url, api_key="EMPTY")
+            client = openai.OpenAI(base_url=base_url, api_key="EMPTY")
 
             try:
                 model_id = client.models.list().data[0].id
@@ -51,7 +52,7 @@ def _initialize_clients(self, base_port: int, max_clients: int) -> None:
 
         logger.info(f"Initialized {len(self.clients)} clients")
 
-    def get_client(self, idx: int) -> OpenAI:
+    def get_client(self, idx: int):
         """Get a client using round-robin load balancing."""
         return self.clients[idx % len(self.clients)]
 
@@ -101,7 +102,7 @@ def _convert_messages(
         return converted_messages, messages
 
     def _generate_response(
-        self, client: OpenAI, messages: List[Dict], **kwargs
+        self, client, messages: List[Dict], **kwargs
     ) -> Optional[str]:
         """
         Generate a response using the OpenAI API.