llama-index

1
from typing import Optional
2

3
from typing_extensions import NotRequired, TypedDict
4

5
from llama_index.legacy.core.llms.types import ChatMessage
6

7
XINFERENCE_MODEL_SIZES = {
8
    "baichuan": 2048,
9
    "baichuan-chat": 2048,
10
    "wizardlm-v1.0": 2048,
11
    "vicuna-v1.3": 2048,
12
    "orca": 2048,
13
    "chatglm": 2048,
14
    "chatglm2": 8192,
15
    "llama-2-chat": 4096,
16
    "llama-2": 4096,
17
}
18

19

20
class ChatCompletionMessage(TypedDict):
21
    role: str
22
    content: Optional[str]
23
    user: NotRequired[str]
24

25

26
def xinference_message_to_history(message: ChatMessage) -> ChatCompletionMessage:
27
    return ChatCompletionMessage(role=message.role, content=message.content)
28

29

30
def xinference_modelname_to_contextsize(modelname: str) -> int:
31
    context_size = XINFERENCE_MODEL_SIZES.get(modelname, None)
32

33
    if context_size is None:
34
        raise ValueError(
35
            f"Unknown model: {modelname}. Please provide a valid OpenAI model name."
36
            "Known models are: " + ", ".join(XINFERENCE_MODEL_SIZES.keys())
37
        )
38

39
    return context_size
40