mirror of
https://github.com/hiyouga/LLaMA-Factory.git
synced 2025-08-02 03:32:50 +08:00
129 lines
3.1 KiB
Python
129 lines
3.1 KiB
Python
import time
|
|
from enum import Enum, unique
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from pydantic import BaseModel, Field
|
|
from typing_extensions import Literal
|
|
|
|
|
|
@unique
|
|
class Role(str, Enum):
|
|
USER = "user"
|
|
ASSISTANT = "assistant"
|
|
SYSTEM = "system"
|
|
FUNCTION = "function"
|
|
TOOL = "tool"
|
|
|
|
|
|
@unique
|
|
class Finish(str, Enum):
|
|
STOP = "stop"
|
|
LENGTH = "length"
|
|
TOOL = "tool_calls"
|
|
|
|
|
|
class ModelCard(BaseModel):
|
|
id: str
|
|
object: Literal["model"] = "model"
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
owned_by: Literal["owner"] = "owner"
|
|
|
|
|
|
class ModelList(BaseModel):
|
|
object: Literal["list"] = "list"
|
|
data: List[ModelCard] = []
|
|
|
|
|
|
class Function(BaseModel):
|
|
name: str
|
|
arguments: str
|
|
|
|
|
|
class FunctionDefinition(BaseModel):
|
|
name: str
|
|
description: str
|
|
parameters: Dict[str, Any]
|
|
|
|
|
|
class FunctionAvailable(BaseModel):
|
|
type: Literal["function", "code_interpreter"] = "function"
|
|
function: Optional[FunctionDefinition] = None
|
|
|
|
|
|
class FunctionCall(BaseModel):
|
|
id: Literal["call_default"] = "call_default"
|
|
type: Literal["function"] = "function"
|
|
function: Function
|
|
|
|
|
|
class ChatMessage(BaseModel):
|
|
role: Role
|
|
content: Optional[str] = None
|
|
tool_calls: Optional[List[FunctionCall]] = None
|
|
|
|
|
|
class ChatCompletionMessage(BaseModel):
|
|
role: Optional[Role] = None
|
|
content: Optional[str] = None
|
|
tool_calls: Optional[List[FunctionCall]] = None
|
|
|
|
|
|
class ChatCompletionRequest(BaseModel):
|
|
model: str
|
|
messages: List[ChatMessage]
|
|
tools: Optional[List[FunctionAvailable]] = None
|
|
do_sample: bool = True
|
|
temperature: Optional[float] = None
|
|
top_p: Optional[float] = None
|
|
n: int = 1
|
|
max_tokens: Optional[int] = None
|
|
stream: bool = False
|
|
|
|
|
|
class ChatCompletionResponseChoice(BaseModel):
|
|
index: int
|
|
message: ChatCompletionMessage
|
|
finish_reason: Finish
|
|
|
|
|
|
class ChatCompletionResponseStreamChoice(BaseModel):
|
|
index: int
|
|
delta: ChatCompletionMessage
|
|
finish_reason: Optional[Finish] = None
|
|
|
|
|
|
class ChatCompletionResponseUsage(BaseModel):
|
|
prompt_tokens: int
|
|
completion_tokens: int
|
|
total_tokens: int
|
|
|
|
|
|
class ChatCompletionResponse(BaseModel):
|
|
id: Literal["chatcmpl-default"] = "chatcmpl-default"
|
|
object: Literal["chat.completion"] = "chat.completion"
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
model: str
|
|
choices: List[ChatCompletionResponseChoice]
|
|
usage: ChatCompletionResponseUsage
|
|
|
|
|
|
class ChatCompletionStreamResponse(BaseModel):
|
|
id: Literal["chatcmpl-default"] = "chatcmpl-default"
|
|
object: Literal["chat.completion.chunk"] = "chat.completion.chunk"
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|
model: str
|
|
choices: List[ChatCompletionResponseStreamChoice]
|
|
|
|
|
|
class ScoreEvaluationRequest(BaseModel):
|
|
model: str
|
|
messages: List[str]
|
|
max_length: Optional[int] = None
|
|
|
|
|
|
class ScoreEvaluationResponse(BaseModel):
|
|
id: Literal["scoreeval-default"] = "scoreeval-default"
|
|
object: Literal["score.evaluation"] = "score.evaluation"
|
|
model: str
|
|
scores: List[float]
|