We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
ENH: increase context length


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `xinference/model/llm/core.py`
Content:
```
1 # Copyright 2022-2023 XProbe Inc.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #      http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 import abc
16 import logging
17 from abc import abstractmethod
18 from typing import TYPE_CHECKING, Any, Iterator, List, Optional, TypedDict, Union
19 
20 from .types import (
21     ChatCompletion,
22     ChatCompletionChunk,
23     ChatCompletionMessage,
24     Completion,
25     CompletionChunk,
26 )
27 
28 if TYPE_CHECKING:
29     from llama_cpp import LogitsProcessorList, StoppingCriteriaList
30 
31     from .. import ModelSpec
32 
33 logger = logging.getLogger(__name__)
34 
35 
36 class StrictTypedDict(TypedDict):
37     def __setitem__(self, key: str, value: Any):  # type: ignore
38         if key not in self.__annotations__:
39             raise KeyError(f"Key '{key}' is not allowed in {self.__class__.__name__}")
40 
41         expected_type = self.__annotations__[key]
42         if not isinstance(value, expected_type):
43             raise TypeError(
44                 f"Value for key '{key}' must be of type '{expected_type.__name__}', "
45                 f"not '{type(value).__name__}'"
46             )
47 
48         super().__setitem__(key, value)
49 
50 
51 class LlamaCppGenerateConfig(StrictTypedDict, total=False):
52     suffix: Optional[str]
53     max_tokens: int
54     temperature: float
55     top_p: float
56     logprobs: Optional[int]
57     echo: bool
58     stop: Optional[Union[str, List[str]]]
59     frequency_penalty: float
60     presence_penalty: float
61     repeat_penalty: float
62     top_k: int
63     stream: bool
64     tfs_z: float
65     mirostat_mode: int
66     mirostat_tau: float
67     mirostat_eta: float
68     model: Optional[str]
69     stopping_criteria: Optional["StoppingCriteriaList"]
70     logits_processor: Optional["LogitsProcessorList"]
71 
72 
73 class LlamaCppModelConfig(StrictTypedDict, total=False):
74     n_ctx: int
75     n_parts: int
76     n_gpu_layers: int
77     seed: int
78     f16_kv: bool
79     logits_all: bool
80     vocab_only: bool
81     use_mmap: bool
82     use_mlock: bool
83     embedding: bool
84     n_threads: Optional[int]
85     n_batch: int
86     last_n_tokens_size: int
87     lora_base: Optional[str]
88     lora_path: Optional[str]
89     low_vram: bool
90     verbose: bool
91 
92 
93 class Model(abc.ABC):
94     def __init__(self, model_uid: str, model_spec: "ModelSpec", *args, **kwargs):
95         self.model_uid = model_uid
96         self.model_spec = model_spec
97 
98     @abstractmethod
99     def load(self):
100         pass
101 
102 
103 class LlamaCppModel(Model):
104     def __init__(
105         self,
106         model_uid: str,
107         model_spec: "ModelSpec",
108         model_path: str,
109         llamacpp_model_config: Optional[LlamaCppModelConfig] = None,
110     ):
111         super().__init__(model_uid, model_spec)
112         self._model_path = model_path
113         self._llamacpp_model_config: LlamaCppModelConfig = self._sanitize_model_config(
114             llamacpp_model_config
115         )
116         self._llm = None
117 
118     @classmethod
119     def _sanitize_model_config(
120         cls, llamacpp_model_config: Optional[LlamaCppModelConfig]
121     ) -> LlamaCppModelConfig:
122         if llamacpp_model_config is None:
123             llamacpp_model_config = LlamaCppModelConfig()
124         return llamacpp_model_config
125 
126     @classmethod
127     def _sanitize_generate_config(
128         cls,
129         generate_config: Optional[LlamaCppGenerateConfig],
130     ) -> LlamaCppGenerateConfig:
131         if generate_config is None:
132             generate_config = LlamaCppGenerateConfig()
133         return generate_config
134 
135     def load(self):
136         from llama_cpp import Llama
137 
138         self._llm = Llama(
139             model_path=self._model_path,
140             verbose=False,
141             **self._llamacpp_model_config,
142         )
143 
144     def generate(
145         self, prompt: str, generate_config: Optional[LlamaCppGenerateConfig] = None
146     ) -> Union[Completion, Iterator[CompletionChunk]]:
147         def generator_wrapper(
148             _prompt: str, _generate_config: LlamaCppGenerateConfig
149         ) -> Iterator[CompletionChunk]:
150             assert self._llm is not None
151             for _completion_chunk in self._llm(prompt=_prompt, **_generate_config):
152                 yield _completion_chunk
153 
154         logger.debug(
155             "Enter generate, prompt: %s, generate config: %s", prompt, generate_config
156         )
157 
158         generate_config = self._sanitize_generate_config(generate_config)
159 
160         stream = generate_config.get("stream", False)
161         if not stream:
162             assert self._llm is not None
163             completion = self._llm(prompt=prompt, **generate_config)
164 
165             return completion
166         else:
167             return generator_wrapper(prompt, generate_config)
168 
169 
170 class LlamaCppChatModel(LlamaCppModel):
171     def __init__(
172         self,
173         model_uid: str,
174         model_spec: "ModelSpec",
175         model_path: str,
176         system_prompt: str,
177         sep: str,
178         user_name: str,
179         assistant_name: str,
180         llamacpp_model_config: Optional[LlamaCppModelConfig] = None,
181     ):
182         super().__init__(model_uid, model_spec, model_path, llamacpp_model_config)
183         self._system_prompt: str = system_prompt
184         self._sep: str = sep
185         self._user_name: str = user_name
186         self._assistant_name: str = assistant_name
187 
188     def _to_prompt(
189         self,
190         prompt: str,
191         system_prompt: str,
192         chat_history: List[ChatCompletionMessage],
193     ):
194         ret = system_prompt
195         for message in chat_history:
196             role = message["role"]
197             content = message["content"]
198             ret += f"{self._sep}{role}: {content}"
199         ret += f"{self._sep}{self._user_name}: {prompt}"
200         ret += f"{self._sep}{self._assistant_name}:"
201         return ret
202 
203     @staticmethod
204     def _convert_chat_completion_chunks_to_chat(
205         chunks: Iterator[CompletionChunk],
206     ) -> Iterator[ChatCompletionChunk]:
207         for i, chunk in enumerate(chunks):
208             if i == 0:
209                 yield {
210                     "id": "chat" + chunk["id"],
211                     "model": chunk["model"],
212                     "created": chunk["created"],
213                     "object": "chat.completion.chunk",
214                     "choices": [
215                         {
216                             "index": 0,
217                             "delta": {
218                                 "role": "assistant",
219                             },
220                             "finish_reason": None,
221                         }
222                     ],
223                 }
224             yield {
225                 "id": "chat" + chunk["id"],
226                 "model": chunk["model"],
227                 "created": chunk["created"],
228                 "object": "chat.completion.chunk",
229                 "choices": [
230                     {
231                         "index": 0,
232                         "delta": {
233                             "content": chunk["choices"][0]["text"],
234                         },
235                         "finish_reason": chunk["choices"][0]["finish_reason"],
236                     }
237                 ],
238             }
239 
240     @staticmethod
241     def _convert_text_completion_to_chat(completion: Completion) -> ChatCompletion:
242         return {
243             "id": "chat" + completion["id"],
244             "object": "chat.completion",
245             "created": completion["created"],
246             "model": completion["model"],
247             "choices": [
248                 {
249                     "index": 0,
250                     "message": {
251                         "role": "assistant",
252                         "content": completion["choices"][0]["text"],
253                     },
254                     "finish_reason": completion["choices"][0]["finish_reason"],
255                 }
256             ],
257             "usage": completion["usage"],
258         }
259 
260     def chat(
261         self,
262         prompt: str,
263         system_prompt: Optional[str] = None,
264         chat_history: Optional[List[ChatCompletionMessage]] = None,
265         generate_config: Optional[LlamaCppGenerateConfig] = None,
266     ) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:
267         system_prompt = system_prompt or self._system_prompt
268         chat_history = chat_history or []
269         full_prompt = self._to_prompt(prompt, system_prompt, chat_history=chat_history)
270 
271         generate_config = self._sanitize_generate_config(generate_config)
272 
273         stream = generate_config.get("stream", False)
274         if stream:
275             it = self.generate(full_prompt, generate_config)
276             assert isinstance(it, Iterator)
277             return self._convert_chat_completion_chunks_to_chat(it)
278         else:
279             c = self.generate(full_prompt, generate_config)
280             assert not isinstance(c, Iterator)
281             return self._convert_text_completion_to_chat(c)
282
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/xinference/model/llm/core.py b/xinference/model/llm/core.py
--- a/xinference/model/llm/core.py
+++ b/xinference/model/llm/core.py
@@ -14,6 +14,7 @@
 
 import abc
 import logging
+import platform
 from abc import abstractmethod
 from typing import TYPE_CHECKING, Any, Iterator, List, Optional, TypedDict, Union
 
@@ -32,6 +33,14 @@
 
 logger = logging.getLogger(__name__)
 
+SIZE_TO_GPU_LAYERS = {
+    3: 26,
+    7: 32,
+    13: 40,
+    30: 60,
+    65: 80,
+}
+
 
 class StrictTypedDict(TypedDict):
     def __setitem__(self, key: str, value: Any):  # type: ignore
@@ -48,7 +57,7 @@
         super().__setitem__(key, value)
 
 
-class LlamaCppGenerateConfig(StrictTypedDict, total=False):
+class LlamaCppGenerateConfig(TypedDict, total=False):
     suffix: Optional[str]
     max_tokens: int
     temperature: float
@@ -70,7 +79,7 @@
     logits_processor: Optional["LogitsProcessorList"]
 
 
-class LlamaCppModelConfig(StrictTypedDict, total=False):
+class LlamaCppModelConfig(TypedDict, total=False):
     n_ctx: int
     n_parts: int
     n_gpu_layers: int
@@ -109,31 +118,53 @@
         llamacpp_model_config: Optional[LlamaCppModelConfig] = None,
     ):
         super().__init__(model_uid, model_spec)
+
+        closest_size = min(
+            SIZE_TO_GPU_LAYERS.keys(),
+            key=lambda x: abs(x - model_spec.model_size_in_billions),
+        )
+        self._gpu_layers = SIZE_TO_GPU_LAYERS[closest_size]
         self._model_path = model_path
         self._llamacpp_model_config: LlamaCppModelConfig = self._sanitize_model_config(
             llamacpp_model_config
         )
         self._llm = None
 
-    @classmethod
     def _sanitize_model_config(
-        cls, llamacpp_model_config: Optional[LlamaCppModelConfig]
+        self, llamacpp_model_config: Optional[LlamaCppModelConfig]
     ) -> LlamaCppModelConfig:
         if llamacpp_model_config is None:
             llamacpp_model_config = LlamaCppModelConfig()
+        if platform.system() == "Windows":
+            context_length = 512
+        else:
+            context_length = 2048
+
+        llamacpp_model_config.setdefault("n_gpu_layers", self._gpu_layers)
+        llamacpp_model_config.setdefault("n_ctx", context_length)
+
         return llamacpp_model_config
 
-    @classmethod
     def _sanitize_generate_config(
-        cls,
-        generate_config: Optional[LlamaCppGenerateConfig],
+        self, generate_config: Optional[LlamaCppGenerateConfig]
     ) -> LlamaCppGenerateConfig:
         if generate_config is None:
             generate_config = LlamaCppGenerateConfig()
+        generate_config["model"] = self.model_uid
         return generate_config
 
     def load(self):
-        from llama_cpp import Llama
+        try:
+            from llama_cpp import Llama
+        except ImportError:
+            error_message = "Failed to import module 'llama_cpp'"
+            installation_guide = [
+                "Please make sure 'llama_cpp' is installed. ",
+                "You can install it by visiting the installation section of the git repo:\n",
+                "https://github.com/abetlen/llama-cpp-python#installation-with-openblas--cublas--clblast--metal",
+            ]
+
+            raise ImportError(f"{error_message}\n\n{''.join(installation_guide)}")
 
         self._llm = Llama(
             model_path=self._model_path,

{"golden_diff": "diff --git a/xinference/model/llm/core.py b/xinference/model/llm/core.py\n--- a/xinference/model/llm/core.py\n+++ b/xinference/model/llm/core.py\n@@ -14,6 +14,7 @@\n \n import abc\n import logging\n+import platform\n from abc import abstractmethod\n from typing import TYPE_CHECKING, Any, Iterator, List, Optional, TypedDict, Union\n \n@@ -32,6 +33,14 @@\n \n logger = logging.getLogger(__name__)\n \n+SIZE_TO_GPU_LAYERS = {\n+    3: 26,\n+    7: 32,\n+    13: 40,\n+    30: 60,\n+    65: 80,\n+}\n+\n \n class StrictTypedDict(TypedDict):\n     def __setitem__(self, key: str, value: Any):  # type: ignore\n@@ -48,7 +57,7 @@\n         super().__setitem__(key, value)\n \n \n-class LlamaCppGenerateConfig(StrictTypedDict, total=False):\n+class LlamaCppGenerateConfig(TypedDict, total=False):\n     suffix: Optional[str]\n     max_tokens: int\n     temperature: float\n@@ -70,7 +79,7 @@\n     logits_processor: Optional[\"LogitsProcessorList\"]\n \n \n-class LlamaCppModelConfig(StrictTypedDict, total=False):\n+class LlamaCppModelConfig(TypedDict, total=False):\n     n_ctx: int\n     n_parts: int\n     n_gpu_layers: int\n@@ -109,31 +118,53 @@\n         llamacpp_model_config: Optional[LlamaCppModelConfig] = None,\n     ):\n         super().__init__(model_uid, model_spec)\n+\n+        closest_size = min(\n+            SIZE_TO_GPU_LAYERS.keys(),\n+            key=lambda x: abs(x - model_spec.model_size_in_billions),\n+        )\n+        self._gpu_layers = SIZE_TO_GPU_LAYERS[closest_size]\n         self._model_path = model_path\n         self._llamacpp_model_config: LlamaCppModelConfig = self._sanitize_model_config(\n             llamacpp_model_config\n         )\n         self._llm = None\n \n-    @classmethod\n     def _sanitize_model_config(\n-        cls, llamacpp_model_config: Optional[LlamaCppModelConfig]\n+        self, llamacpp_model_config: Optional[LlamaCppModelConfig]\n     ) -> LlamaCppModelConfig:\n         if llamacpp_model_config is None:\n             llamacpp_model_config = LlamaCppModelConfig()\n+        if platform.system() == \"Windows\":\n+            context_length = 512\n+        else:\n+            context_length = 2048\n+\n+        llamacpp_model_config.setdefault(\"n_gpu_layers\", self._gpu_layers)\n+        llamacpp_model_config.setdefault(\"n_ctx\", context_length)\n+\n         return llamacpp_model_config\n \n-    @classmethod\n     def _sanitize_generate_config(\n-        cls,\n-        generate_config: Optional[LlamaCppGenerateConfig],\n+        self, generate_config: Optional[LlamaCppGenerateConfig]\n     ) -> LlamaCppGenerateConfig:\n         if generate_config is None:\n             generate_config = LlamaCppGenerateConfig()\n+        generate_config[\"model\"] = self.model_uid\n         return generate_config\n \n     def load(self):\n-        from llama_cpp import Llama\n+        try:\n+            from llama_cpp import Llama\n+        except ImportError:\n+            error_message = \"Failed to import module 'llama_cpp'\"\n+            installation_guide = [\n+                \"Please make sure 'llama_cpp' is installed. \",\n+                \"You can install it by visiting the installation section of the git repo:\\n\",\n+                \"https://github.com/abetlen/llama-cpp-python#installation-with-openblas--cublas--clblast--metal\",\n+            ]\n+\n+            raise ImportError(f\"{error_message}\\n\\n{''.join(installation_guide)}\")\n \n         self._llm = Llama(\n             model_path=self._model_path,\n", "issue": "ENH: increase context length\n\n", "before_files": [{"content": "# Copyright 2022-2023 XProbe Inc.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport abc\nimport logging\nfrom abc import abstractmethod\nfrom typing import TYPE_CHECKING, Any, Iterator, List, Optional, TypedDict, Union\n\nfrom .types import (\n    ChatCompletion,\n    ChatCompletionChunk,\n    ChatCompletionMessage,\n    Completion,\n    CompletionChunk,\n)\n\nif TYPE_CHECKING:\n    from llama_cpp import LogitsProcessorList, StoppingCriteriaList\n\n    from .. import ModelSpec\n\nlogger = logging.getLogger(__name__)\n\n\nclass StrictTypedDict(TypedDict):\n    def __setitem__(self, key: str, value: Any):  # type: ignore\n        if key not in self.__annotations__:\n            raise KeyError(f\"Key '{key}' is not allowed in {self.__class__.__name__}\")\n\n        expected_type = self.__annotations__[key]\n        if not isinstance(value, expected_type):\n            raise TypeError(\n                f\"Value for key '{key}' must be of type '{expected_type.__name__}', \"\n                f\"not '{type(value).__name__}'\"\n            )\n\n        super().__setitem__(key, value)\n\n\nclass LlamaCppGenerateConfig(StrictTypedDict, total=False):\n    suffix: Optional[str]\n    max_tokens: int\n    temperature: float\n    top_p: float\n    logprobs: Optional[int]\n    echo: bool\n    stop: Optional[Union[str, List[str]]]\n    frequency_penalty: float\n    presence_penalty: float\n    repeat_penalty: float\n    top_k: int\n    stream: bool\n    tfs_z: float\n    mirostat_mode: int\n    mirostat_tau: float\n    mirostat_eta: float\n    model: Optional[str]\n    stopping_criteria: Optional[\"StoppingCriteriaList\"]\n    logits_processor: Optional[\"LogitsProcessorList\"]\n\n\nclass LlamaCppModelConfig(StrictTypedDict, total=False):\n    n_ctx: int\n    n_parts: int\n    n_gpu_layers: int\n    seed: int\n    f16_kv: bool\n    logits_all: bool\n    vocab_only: bool\n    use_mmap: bool\n    use_mlock: bool\n    embedding: bool\n    n_threads: Optional[int]\n    n_batch: int\n    last_n_tokens_size: int\n    lora_base: Optional[str]\n    lora_path: Optional[str]\n    low_vram: bool\n    verbose: bool\n\n\nclass Model(abc.ABC):\n    def __init__(self, model_uid: str, model_spec: \"ModelSpec\", *args, **kwargs):\n        self.model_uid = model_uid\n        self.model_spec = model_spec\n\n    @abstractmethod\n    def load(self):\n        pass\n\n\nclass LlamaCppModel(Model):\n    def __init__(\n        self,\n        model_uid: str,\n        model_spec: \"ModelSpec\",\n        model_path: str,\n        llamacpp_model_config: Optional[LlamaCppModelConfig] = None,\n    ):\n        super().__init__(model_uid, model_spec)\n        self._model_path = model_path\n        self._llamacpp_model_config: LlamaCppModelConfig = self._sanitize_model_config(\n            llamacpp_model_config\n        )\n        self._llm = None\n\n    @classmethod\n    def _sanitize_model_config(\n        cls, llamacpp_model_config: Optional[LlamaCppModelConfig]\n    ) -> LlamaCppModelConfig:\n        if llamacpp_model_config is None:\n            llamacpp_model_config = LlamaCppModelConfig()\n        return llamacpp_model_config\n\n    @classmethod\n    def _sanitize_generate_config(\n        cls,\n        generate_config: Optional[LlamaCppGenerateConfig],\n    ) -> LlamaCppGenerateConfig:\n        if generate_config is None:\n            generate_config = LlamaCppGenerateConfig()\n        return generate_config\n\n    def load(self):\n        from llama_cpp import Llama\n\n        self._llm = Llama(\n            model_path=self._model_path,\n            verbose=False,\n            **self._llamacpp_model_config,\n        )\n\n    def generate(\n        self, prompt: str, generate_config: Optional[LlamaCppGenerateConfig] = None\n    ) -> Union[Completion, Iterator[CompletionChunk]]:\n        def generator_wrapper(\n            _prompt: str, _generate_config: LlamaCppGenerateConfig\n        ) -> Iterator[CompletionChunk]:\n            assert self._llm is not None\n            for _completion_chunk in self._llm(prompt=_prompt, **_generate_config):\n                yield _completion_chunk\n\n        logger.debug(\n            \"Enter generate, prompt: %s, generate config: %s\", prompt, generate_config\n        )\n\n        generate_config = self._sanitize_generate_config(generate_config)\n\n        stream = generate_config.get(\"stream\", False)\n        if not stream:\n            assert self._llm is not None\n            completion = self._llm(prompt=prompt, **generate_config)\n\n            return completion\n        else:\n            return generator_wrapper(prompt, generate_config)\n\n\nclass LlamaCppChatModel(LlamaCppModel):\n    def __init__(\n        self,\n        model_uid: str,\n        model_spec: \"ModelSpec\",\n        model_path: str,\n        system_prompt: str,\n        sep: str,\n        user_name: str,\n        assistant_name: str,\n        llamacpp_model_config: Optional[LlamaCppModelConfig] = None,\n    ):\n        super().__init__(model_uid, model_spec, model_path, llamacpp_model_config)\n        self._system_prompt: str = system_prompt\n        self._sep: str = sep\n        self._user_name: str = user_name\n        self._assistant_name: str = assistant_name\n\n    def _to_prompt(\n        self,\n        prompt: str,\n        system_prompt: str,\n        chat_history: List[ChatCompletionMessage],\n    ):\n        ret = system_prompt\n        for message in chat_history:\n            role = message[\"role\"]\n            content = message[\"content\"]\n            ret += f\"{self._sep}{role}: {content}\"\n        ret += f\"{self._sep}{self._user_name}: {prompt}\"\n        ret += f\"{self._sep}{self._assistant_name}:\"\n        return ret\n\n    @staticmethod\n    def _convert_chat_completion_chunks_to_chat(\n        chunks: Iterator[CompletionChunk],\n    ) -> Iterator[ChatCompletionChunk]:\n        for i, chunk in enumerate(chunks):\n            if i == 0:\n                yield {\n                    \"id\": \"chat\" + chunk[\"id\"],\n                    \"model\": chunk[\"model\"],\n                    \"created\": chunk[\"created\"],\n                    \"object\": \"chat.completion.chunk\",\n                    \"choices\": [\n                        {\n                            \"index\": 0,\n                            \"delta\": {\n                                \"role\": \"assistant\",\n                            },\n                            \"finish_reason\": None,\n                        }\n                    ],\n                }\n            yield {\n                \"id\": \"chat\" + chunk[\"id\"],\n                \"model\": chunk[\"model\"],\n                \"created\": chunk[\"created\"],\n                \"object\": \"chat.completion.chunk\",\n                \"choices\": [\n                    {\n                        \"index\": 0,\n                        \"delta\": {\n                            \"content\": chunk[\"choices\"][0][\"text\"],\n                        },\n                        \"finish_reason\": chunk[\"choices\"][0][\"finish_reason\"],\n                    }\n                ],\n            }\n\n    @staticmethod\n    def _convert_text_completion_to_chat(completion: Completion) -> ChatCompletion:\n        return {\n            \"id\": \"chat\" + completion[\"id\"],\n            \"object\": \"chat.completion\",\n            \"created\": completion[\"created\"],\n            \"model\": completion[\"model\"],\n            \"choices\": [\n                {\n                    \"index\": 0,\n                    \"message\": {\n                        \"role\": \"assistant\",\n                        \"content\": completion[\"choices\"][0][\"text\"],\n                    },\n                    \"finish_reason\": completion[\"choices\"][0][\"finish_reason\"],\n                }\n            ],\n            \"usage\": completion[\"usage\"],\n        }\n\n    def chat(\n        self,\n        prompt: str,\n        system_prompt: Optional[str] = None,\n        chat_history: Optional[List[ChatCompletionMessage]] = None,\n        generate_config: Optional[LlamaCppGenerateConfig] = None,\n    ) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:\n        system_prompt = system_prompt or self._system_prompt\n        chat_history = chat_history or []\n        full_prompt = self._to_prompt(prompt, system_prompt, chat_history=chat_history)\n\n        generate_config = self._sanitize_generate_config(generate_config)\n\n        stream = generate_config.get(\"stream\", False)\n        if stream:\n            it = self.generate(full_prompt, generate_config)\n            assert isinstance(it, Iterator)\n            return self._convert_chat_completion_chunks_to_chat(it)\n        else:\n            c = self.generate(full_prompt, generate_config)\n            assert not isinstance(c, Iterator)\n            return self._convert_text_completion_to_chat(c)\n", "path": "xinference/model/llm/core.py"}], "after_files": [{"content": "# Copyright 2022-2023 XProbe Inc.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport abc\nimport logging\nimport platform\nfrom abc import abstractmethod\nfrom typing import TYPE_CHECKING, Any, Iterator, List, Optional, TypedDict, Union\n\nfrom .types import (\n    ChatCompletion,\n    ChatCompletionChunk,\n    ChatCompletionMessage,\n    Completion,\n    CompletionChunk,\n)\n\nif TYPE_CHECKING:\n    from llama_cpp import LogitsProcessorList, StoppingCriteriaList\n\n    from .. import ModelSpec\n\nlogger = logging.getLogger(__name__)\n\nSIZE_TO_GPU_LAYERS = {\n    3: 26,\n    7: 32,\n    13: 40,\n    30: 60,\n    65: 80,\n}\n\n\nclass StrictTypedDict(TypedDict):\n    def __setitem__(self, key: str, value: Any):  # type: ignore\n        if key not in self.__annotations__:\n            raise KeyError(f\"Key '{key}' is not allowed in {self.__class__.__name__}\")\n\n        expected_type = self.__annotations__[key]\n        if not isinstance(value, expected_type):\n            raise TypeError(\n                f\"Value for key '{key}' must be of type '{expected_type.__name__}', \"\n                f\"not '{type(value).__name__}'\"\n            )\n\n        super().__setitem__(key, value)\n\n\nclass LlamaCppGenerateConfig(TypedDict, total=False):\n    suffix: Optional[str]\n    max_tokens: int\n    temperature: float\n    top_p: float\n    logprobs: Optional[int]\n    echo: bool\n    stop: Optional[Union[str, List[str]]]\n    frequency_penalty: float\n    presence_penalty: float\n    repeat_penalty: float\n    top_k: int\n    stream: bool\n    tfs_z: float\n    mirostat_mode: int\n    mirostat_tau: float\n    mirostat_eta: float\n    model: Optional[str]\n    stopping_criteria: Optional[\"StoppingCriteriaList\"]\n    logits_processor: Optional[\"LogitsProcessorList\"]\n\n\nclass LlamaCppModelConfig(TypedDict, total=False):\n    n_ctx: int\n    n_parts: int\n    n_gpu_layers: int\n    seed: int\n    f16_kv: bool\n    logits_all: bool\n    vocab_only: bool\n    use_mmap: bool\n    use_mlock: bool\n    embedding: bool\n    n_threads: Optional[int]\n    n_batch: int\n    last_n_tokens_size: int\n    lora_base: Optional[str]\n    lora_path: Optional[str]\n    low_vram: bool\n    verbose: bool\n\n\nclass Model(abc.ABC):\n    def __init__(self, model_uid: str, model_spec: \"ModelSpec\", *args, **kwargs):\n        self.model_uid = model_uid\n        self.model_spec = model_spec\n\n    @abstractmethod\n    def load(self):\n        pass\n\n\nclass LlamaCppModel(Model):\n    def __init__(\n        self,\n        model_uid: str,\n        model_spec: \"ModelSpec\",\n        model_path: str,\n        llamacpp_model_config: Optional[LlamaCppModelConfig] = None,\n    ):\n        super().__init__(model_uid, model_spec)\n\n        closest_size = min(\n            SIZE_TO_GPU_LAYERS.keys(),\n            key=lambda x: abs(x - model_spec.model_size_in_billions),\n        )\n        self._gpu_layers = SIZE_TO_GPU_LAYERS[closest_size]\n        self._model_path = model_path\n        self._llamacpp_model_config: LlamaCppModelConfig = self._sanitize_model_config(\n            llamacpp_model_config\n        )\n        self._llm = None\n\n    def _sanitize_model_config(\n        self, llamacpp_model_config: Optional[LlamaCppModelConfig]\n    ) -> LlamaCppModelConfig:\n        if llamacpp_model_config is None:\n            llamacpp_model_config = LlamaCppModelConfig()\n        if platform.system() == \"Windows\":\n            context_length = 512\n        else:\n            context_length = 2048\n\n        llamacpp_model_config.setdefault(\"n_gpu_layers\", self._gpu_layers)\n        llamacpp_model_config.setdefault(\"n_ctx\", context_length)\n\n        return llamacpp_model_config\n\n    def _sanitize_generate_config(\n        self, generate_config: Optional[LlamaCppGenerateConfig]\n    ) -> LlamaCppGenerateConfig:\n        if generate_config is None:\n            generate_config = LlamaCppGenerateConfig()\n        generate_config[\"model\"] = self.model_uid\n        return generate_config\n\n    def load(self):\n        try:\n            from llama_cpp import Llama\n        except ImportError:\n            error_message = \"Failed to import module 'llama_cpp'\"\n            installation_guide = [\n                \"Please make sure 'llama_cpp' is installed. \",\n                \"You can install it by visiting the installation section of the git repo:\\n\",\n                \"https://github.com/abetlen/llama-cpp-python#installation-with-openblas--cublas--clblast--metal\",\n            ]\n\n            raise ImportError(f\"{error_message}\\n\\n{''.join(installation_guide)}\")\n\n        self._llm = Llama(\n            model_path=self._model_path,\n            verbose=False,\n            **self._llamacpp_model_config,\n        )\n\n    def generate(\n        self, prompt: str, generate_config: Optional[LlamaCppGenerateConfig] = None\n    ) -> Union[Completion, Iterator[CompletionChunk]]:\n        def generator_wrapper(\n            _prompt: str, _generate_config: LlamaCppGenerateConfig\n        ) -> Iterator[CompletionChunk]:\n            assert self._llm is not None\n            for _completion_chunk in self._llm(prompt=_prompt, **_generate_config):\n                yield _completion_chunk\n\n        logger.debug(\n            \"Enter generate, prompt: %s, generate config: %s\", prompt, generate_config\n        )\n\n        generate_config = self._sanitize_generate_config(generate_config)\n\n        stream = generate_config.get(\"stream\", False)\n        if not stream:\n            assert self._llm is not None\n            completion = self._llm(prompt=prompt, **generate_config)\n\n            return completion\n        else:\n            return generator_wrapper(prompt, generate_config)\n\n\nclass LlamaCppChatModel(LlamaCppModel):\n    def __init__(\n        self,\n        model_uid: str,\n        model_spec: \"ModelSpec\",\n        model_path: str,\n        system_prompt: str,\n        sep: str,\n        user_name: str,\n        assistant_name: str,\n        llamacpp_model_config: Optional[LlamaCppModelConfig] = None,\n    ):\n        super().__init__(model_uid, model_spec, model_path, llamacpp_model_config)\n        self._system_prompt: str = system_prompt\n        self._sep: str = sep\n        self._user_name: str = user_name\n        self._assistant_name: str = assistant_name\n\n    def _to_prompt(\n        self,\n        prompt: str,\n        system_prompt: str,\n        chat_history: List[ChatCompletionMessage],\n    ):\n        ret = system_prompt\n        for message in chat_history:\n            role = message[\"role\"]\n            content = message[\"content\"]\n            ret += f\"{self._sep}{role}: {content}\"\n        ret += f\"{self._sep}{self._user_name}: {prompt}\"\n        ret += f\"{self._sep}{self._assistant_name}:\"\n        return ret\n\n    @staticmethod\n    def _convert_chat_completion_chunks_to_chat(\n        chunks: Iterator[CompletionChunk],\n    ) -> Iterator[ChatCompletionChunk]:\n        for i, chunk in enumerate(chunks):\n            if i == 0:\n                yield {\n                    \"id\": \"chat\" + chunk[\"id\"],\n                    \"model\": chunk[\"model\"],\n                    \"created\": chunk[\"created\"],\n                    \"object\": \"chat.completion.chunk\",\n                    \"choices\": [\n                        {\n                            \"index\": 0,\n                            \"delta\": {\n                                \"role\": \"assistant\",\n                            },\n                            \"finish_reason\": None,\n                        }\n                    ],\n                }\n            yield {\n                \"id\": \"chat\" + chunk[\"id\"],\n                \"model\": chunk[\"model\"],\n                \"created\": chunk[\"created\"],\n                \"object\": \"chat.completion.chunk\",\n                \"choices\": [\n                    {\n                        \"index\": 0,\n                        \"delta\": {\n                            \"content\": chunk[\"choices\"][0][\"text\"],\n                        },\n                        \"finish_reason\": chunk[\"choices\"][0][\"finish_reason\"],\n                    }\n                ],\n            }\n\n    @staticmethod\n    def _convert_text_completion_to_chat(completion: Completion) -> ChatCompletion:\n        return {\n            \"id\": \"chat\" + completion[\"id\"],\n            \"object\": \"chat.completion\",\n            \"created\": completion[\"created\"],\n            \"model\": completion[\"model\"],\n            \"choices\": [\n                {\n                    \"index\": 0,\n                    \"message\": {\n                        \"role\": \"assistant\",\n                        \"content\": completion[\"choices\"][0][\"text\"],\n                    },\n                    \"finish_reason\": completion[\"choices\"][0][\"finish_reason\"],\n                }\n            ],\n            \"usage\": completion[\"usage\"],\n        }\n\n    def chat(\n        self,\n        prompt: str,\n        system_prompt: Optional[str] = None,\n        chat_history: Optional[List[ChatCompletionMessage]] = None,\n        generate_config: Optional[LlamaCppGenerateConfig] = None,\n    ) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:\n        system_prompt = system_prompt or self._system_prompt\n        chat_history = chat_history or []\n        full_prompt = self._to_prompt(prompt, system_prompt, chat_history=chat_history)\n\n        generate_config = self._sanitize_generate_config(generate_config)\n\n        stream = generate_config.get(\"stream\", False)\n        if stream:\n            it = self.generate(full_prompt, generate_config)\n            assert isinstance(it, Iterator)\n            return self._convert_chat_completion_chunks_to_chat(it)\n        else:\n            c = self.generate(full_prompt, generate_config)\n            assert not isinstance(c, Iterator)\n            return self._convert_text_completion_to_chat(c)\n", "path": "xinference/model/llm/core.py"}]}