Set up OpenAI API Python Dataclasses

Author: Varun Puri

api/api.py

@@ -0,0 +1,218 @@
+import time
+import uuid
+from abc import ABC
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+from build.utils import device_sync
+
+from generate import Generator, GeneratorArgs
+
+
+@dataclass
+class AbstractMessageType(ABC):
+    role: str
+    content: Optional[str] = None
+
+
+@dataclass
+class SystemMessageType(AbstractMessageType):
+    role: str = "system"
+    name: Optional[str] = None
+
+
+@dataclass
+class UserMessageType(AbstractMessageType):
+    role: str = "user"
+
+
+@dataclass
+class ToolCall:
+    @dataclass
+    class ToolCallFunction:
+        name: str
+        arguments: str
+
+    id: str
+    type: str
+    function: ToolCallFunction
+
+
+@dataclass
+class AssistantMessageType(AbstractMessageType):
+    role: str = "assistant"
+    name: Optional[str] = None
+    tool_calls: Optional[List[ToolCall]] = None
+
+
+@dataclass
+class ToolMessage(AbstractMessageType):
+    role: str = "tool"
+    tool_call_id: Optional[str] = None
+
+
+@dataclass
+class CompletionRequest:
+    model: str
+    prompt: str
+    messages: Optional[List[AbstractMessageType]]
+    frequency_penalty: float = 0.0
+    temperature: float = 0.0
+    stop: Optional[List[str]] = None
+    echo: bool = False
+    frequency_penalty: float = 0.0
+    guided_decode_json_schema: str = None
+    guided_decode_json_schema_path: str = None
+    n: int = 1
+    presence_penalty: float = 0
+    logit_bias: Optional[Dict[str, float]] = None
+    logprobs: Optional[bool] = None
+    top_logprobs: Optional[int] = None
+    max_tokens: Optional[int] = None
+
+
+@dataclass
+class CompletionChoice:
+    finish_reason: str
+    index: int
+    message: AssistantMessageType
+    logprobs: Optional[List[Any]]
+
+
+@dataclass
+class UsageStats:
+    completion_tokens: int
+    prompt_tokens: int
+    total_tokens: int
+
+
+@dataclass
+class CompletionResponse:
+
+    id: str
+    choices: List[CompletionChoice]
+    created: int
+    model: str
+    system_fingerprint: str
+    usage: UsageStats
+    object: str = "chat.completion"
+    service_tier: Optional[str] = None
+
+
+@dataclass
+class ChunkDelta:
+    tool_calls: Optional[List[ToolCall]]
+    role: Optional[str]
+    content: Optional[str]
+
+
+@dataclass
+class CompletionChoiceChunk:
+    delta: ChunkDelta
+    index: int
+    finish_reason: Optional[str] = None
+    logprobs: Optional[List[Any]] = None
+
+
+@dataclass
+class CompletionResponseChunk:
+
+    id: str
+    choices: List[CompletionChoiceChunk]
+    created: int
+    model: str
+    system_fingerprint: str
+    object: str = "chat.completion.chunk"
+    service_tier: Optional[str] = None
+    usage: Optional[UsageStats] = None
+
+
+class OpenAIAPIGenerator(Generator):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.start_pos = 0
+        self.max_seq_length = (
+            self.model.config.max_seq_length
+            + self.speculative_builder_args.speculate_k
+            + 1
+            if self.draft_model is not None
+            else self.model.config.max_seq_length
+        )
+
+    def completion(self, completion_request: CompletionRequest):
+        device_sync(device=self.builder_args.device)
+
+        id = str(uuid.uuid4())
+        idx = 0
+        buffer = []
+        encoded = self.encode_tokens(
+            completion_request.prompt, bos=True, device=self.builder_args.device
+        )
+        generator_args = GeneratorArgs(
+            completion_request.prompt,
+            encoded_prompt=encoded,
+            chat_mode=False,
+        )
+
+        def callback(x, *, done_generating=False):
+            return self._callback(
+                x,
+                buffer=buffer,
+                done_generating=done_generating,
+            )
+
+        for y, _ in self.generate(
+            self.model,
+            encoded,
+            generator_args.max_new_tokens,
+            draft_model=self.draft_model,
+            speculate_k=generator_args.speculate_k,
+            chat_mode=generator_args.chat_mode,
+            callback=callback,
+            temperature=generator_args.temperature,
+            top_k=generator_args.top_k,
+            sequential_prefill=generator_args.sequential_prefill,
+            start_pos=self.start_pos,
+            max_seq_length=self.max_seq_length,
+        ):
+            content = "".join(
+                self.tokenizer.decode([self.tokenizer.encode(".")[0]] + y.tolist())[1:]
+            )
+            chunk_delta = ChunkDelta(
+                role="assistant",
+                content=content,
+                tool_calls=None,
+            )
+            choice_chunk = CompletionChoiceChunk(
+                delta=chunk_delta,
+                index=idx,
+            )
+            chunk_response = CompletionResponseChunk(
+                id=str(id),
+                choices=[choice_chunk],
+                created=int(time.time()),
+                model=completion_request.model,
+                system_fingerprint=uuid.UUID(int=uuid.getnode()),
+            )
+            yield chunk_response
+            self.start_pos += y.size(0)
+            idx += 1
+
+        end_chunk = CompletionChoiceChunk(ChunkDelta(None, None, None), idx, "eos")
+
+        yield CompletionResponseChunk(
+            id=str(id),
+            choices=[end_chunk],
+            created=int(time.time()),
+            model=completion_request.model,
+            system_fingerprint=uuid.UUID(int=uuid.getnode()),
+        )
+
+    def _callback(self, x, *, buffer, done_generating):
+        period_id = self.tokenizer.encode(".")[0]
+        buffer.append(self.tokenizer.decode([period_id] + x.tolist())[1:])
+        if (
+            self.is_llama3_model
+            and x.item() == self.tokenizer.special_tokens["<|eot_id|>"]
+        ):
+            buffer = buffer[:-1]  # drop the eot_id from the output buffer