chore: import upstream snapshot with attribution

research/BGE_Coder/data_generation/llm.py

@@ -0,0 +1,134 @@
+import os
+import time
+import openai
+import random
+import tiktoken
+import threading
+from openai import OpenAI, AzureOpenAI
+from typing import Tuple
+
+
+class LLM:
+    def __init__(
+        self,
+        model: str="Qwen2-5-Coder-32B-Instruct",
+        model_type: str = "open-source",
+        port: int = 8000,
+    ):
+        if model_type == "open-source":
+            self.client = OpenAI(
+                api_key="EMPTY",
+                base_url=f"http://localhost:{port}/v1/"
+            )
+        elif model_type == "azure":
+            self.client = AzureOpenAI(
+                api_key=os.getenv("OPENAI_API_KEY"),
+                api_version=os.getenv("AZURE_API_VERSION", "2024-02-01"),
+                azure_endpoint=os.getenv("AZURE_ENDPOINT"),
+                azure_deployment=os.getenv("OPENAI_DEPLOYMENT_NAME", 'gpt-35-turbo')
+            )
+        elif model_type == "openai":
+            self.client = OpenAI(
+                api_key=os.getenv("OPENAI_API_KEY"),
+                base_url=os.getenv("OPENAI_BASE_URL", None)
+            )
+        else:
+            raise ValueError("model_type must be one of ['open-source', 'azure', 'openai']")
+        
+        self.model = model
+        self.tokenizer = tiktoken.get_encoding("o200k_base")
+    
+    def split_text(self, text: str, anchor_points: Tuple[float, float] = (0.4, 0.7)):
+        token_ids = self.tokenizer.encode(text)
+        anchor_point = random.uniform(anchor_points[0], anchor_points[1])
+        split_index = int(len(token_ids) * anchor_point)
+        return self.tokenizer.decode(token_ids[:split_index]), self.tokenizer.decode(token_ids[split_index:])
+    
+    def chat(
+        self,
+        prompt: str,
+        max_tokens: int = 8192,
+        logit_bais: dict = None,
+        n: int = 1,
+        temperature: float = 1.0,
+        top_p: float = 0.6,
+        repetition_penalty: float = 1.0,
+        remove_thinking: bool = True,
+        timeout: int = 90,
+    ):
+        endure_time = 0
+        endure_time_limit = timeout * 2
+        
+        def create_completion(results):
+            try:
+                completion = self.client.chat.completions.create(
+                    model=self.model,
+                    messages=[{"role": "user", "content": prompt}],
+                    max_tokens=max_tokens,
+                    logit_bias=logit_bais if logit_bais is not None else {},
+                    n=n,
+                    temperature=temperature,
+                    top_p=top_p,
+                    extra_body={'repetition_penalty': repetition_penalty},
+                    timeout=timeout,
+                )
+                results["content"] = [x.message.content for x in completion.choices[:n]]
+            except openai.BadRequestError as e:
+                # The response was filtered due to the prompt triggering Azure OpenAI's content management policy.
+                results["content"] = [None for _ in range(n)]
+            except openai.APIConnectionError as e:
+                results["error"] = f'APIConnectionError({e})'
+            except openai.RateLimitError as e:
+                results["error"] = f'RateLimitError({e})'
+            except Exception as e:
+                results["error"] = f"Error: {e}"
+        
+        while True:
+            results = {"content": None, "error": None}
+            completion_thread = threading.Thread(target=create_completion, args=(results,))
+            completion_thread.start()
+            
+            start_time = time.time()
+            while completion_thread.is_alive():
+                elapsed_time = time.time() - start_time
+                if elapsed_time > endure_time_limit:
+                    print("Completion timeout exceeded. Aborting...")
+                    return [None for _ in range(n)]
+                time.sleep(1)
+            
+            # If an error occurred during result processing
+            if results["error"]:
+                if endure_time >= endure_time_limit:
+                    print(f'{results["error"]} - Skip this prompt.')
+                    return [None for _ in range(n)]
+                print(f"{results['error']} - Waiting for 5 seconds...")
+                endure_time += 5
+                time.sleep(5)
+                continue
+            
+            content_list = results["content"]
+            if remove_thinking:
+                content_list = [x.split('</think>')[-1].strip('\n').strip() if x is not None else None for x in content_list]
+            return content_list
+
+
+if __name__ == "__main__":
+    llm = LLM(
+        model="gpt-4o-mini-2024-07-18",
+        model_type="openai"
+    )
+
+    prompt = "hello, who are you?"
+    response = llm.chat(prompt)[0]
+    print(response)
+
+
+if __name__ == "__main__":
+    llm = LLM(
+        model="gpt-4o-mini-2024-07-18",
+        model_type="openai"
+    )
+
+    prompt = "hello, who are you?"
+    response = llm.chat(prompt)[0]
+    print(response)