We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
FEAT: PyTorch model embeddings


--- 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/pytorch/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 logging
16 from typing import Iterator, List, Optional, TypedDict, Union
17 
18 from ....constants import XINFERENCE_CACHE_DIR
19 from ....types import (
20     ChatCompletion,
21     ChatCompletionChunk,
22     ChatCompletionMessage,
23     Completion,
24     CompletionChunk,
25     Embedding,
26 )
27 from ..core import LLM
28 from ..llm_family import LLMFamilyV1, LLMSpecV1
29 from ..utils import ChatModelMixin
30 
31 logger = logging.getLogger(__name__)
32 
33 
34 class PytorchGenerateConfig(TypedDict, total=False):
35     temperature: float
36     repetition_penalty: float
37     top_p: float
38     top_k: int
39     stream: bool
40     max_new_tokens: int
41     echo: bool
42     stop: Optional[Union[str, List[str]]]
43     stop_token_ids: Optional[Union[int, List[int]]]
44     stream_interval: int
45     model: Optional[str]
46 
47 
48 class PytorchModelConfig(TypedDict, total=False):
49     revision: str
50     device: str
51     gpus: Optional[str]
52     num_gpus: int
53     max_gpu_memory: str
54     gptq_ckpt: Optional[str]
55     gptq_wbits: int
56     gptq_groupsize: int
57     gptq_act_order: bool
58 
59 
60 class PytorchModel(LLM):
61     def __init__(
62         self,
63         model_uid: str,
64         model_family: "LLMFamilyV1",
65         model_spec: "LLMSpecV1",
66         quantization: str,
67         model_path: str,
68         pytorch_model_config: Optional[PytorchModelConfig] = None,
69     ):
70         super().__init__(model_uid, model_family, model_spec, quantization, model_path)
71         self._use_fast_tokenizer = True
72         self._pytorch_model_config: PytorchModelConfig = self._sanitize_model_config(
73             pytorch_model_config
74         )
75 
76     def _sanitize_model_config(
77         self, pytorch_model_config: Optional[PytorchModelConfig]
78     ) -> PytorchModelConfig:
79         if pytorch_model_config is None:
80             pytorch_model_config = PytorchModelConfig()
81         pytorch_model_config.setdefault("revision", "main")
82         pytorch_model_config.setdefault("gpus", None)
83         pytorch_model_config.setdefault("num_gpus", 1)
84         pytorch_model_config.setdefault("gptq_ckpt", None)
85         pytorch_model_config.setdefault("gptq_wbits", 16)
86         pytorch_model_config.setdefault("gptq_groupsize", -1)
87         pytorch_model_config.setdefault("gptq_act_order", False)
88         if self._is_darwin_and_apple_silicon():
89             pytorch_model_config.setdefault("device", "mps")
90         else:
91             pytorch_model_config.setdefault("device", "cuda")
92         return pytorch_model_config
93 
94     def _sanitize_generate_config(
95         self,
96         pytorch_generate_config: Optional[PytorchGenerateConfig],
97     ) -> PytorchGenerateConfig:
98         if pytorch_generate_config is None:
99             pytorch_generate_config = PytorchGenerateConfig()
100         pytorch_generate_config.setdefault("temperature", 0.7)
101         pytorch_generate_config.setdefault("repetition_penalty", 1.0)
102         pytorch_generate_config.setdefault("max_new_tokens", 512)
103         pytorch_generate_config.setdefault("stream_interval", 2)
104         pytorch_generate_config["model"] = self.model_uid
105         return pytorch_generate_config
106 
107     def _load_model(self, kwargs: dict):
108         try:
109             from transformers import AutoModelForCausalLM, AutoTokenizer
110         except ImportError:
111             error_message = "Failed to import module 'transformers'"
112             installation_guide = [
113                 "Please make sure 'transformers' is installed. ",
114                 "You can install it by `pip install transformers`\n",
115             ]
116 
117             raise ImportError(f"{error_message}\n\n{''.join(installation_guide)}")
118 
119         tokenizer = AutoTokenizer.from_pretrained(
120             self.model_path,
121             use_fast=self._use_fast_tokenizer,
122             revision=kwargs["revision"],
123             cache_dir=XINFERENCE_CACHE_DIR,
124         )
125         model = AutoModelForCausalLM.from_pretrained(
126             self.model_path,
127             low_cpu_mem_usage=True,
128             cache_dir=XINFERENCE_CACHE_DIR,
129             **kwargs,
130         )
131         return model, tokenizer
132 
133     def load(self):
134         try:
135             import torch
136         except ImportError:
137             raise ImportError(
138                 f"Failed to import module 'torch'. Please make sure 'torch' is installed.\n\n"
139             )
140         from .compression import load_compress_model
141 
142         quantization = self.quantization
143         num_gpus = self._pytorch_model_config.get("num_gpus", 1)
144         if self._is_darwin_and_apple_silicon():
145             device = self._pytorch_model_config.get("device", "mps")
146         else:
147             device = self._pytorch_model_config.get("device", "cuda")
148 
149         if device == "cpu":
150             kwargs = {"torch_dtype": torch.float32}
151         elif device == "cuda":
152             kwargs = {"torch_dtype": torch.float16}
153         elif device == "mps":
154             kwargs = {"torch_dtype": torch.float16}
155         else:
156             raise ValueError(f"Device {device} is not supported in temporary")
157         kwargs["revision"] = self._pytorch_model_config.get("revision", "main")
158 
159         if quantization != "none":
160             if device == "cuda" and self._is_linux():
161                 kwargs["device_map"] = "auto"
162                 if quantization == "4-bit":
163                     kwargs["load_in_4bit"] = True
164                 elif quantization == "8-bit":
165                     kwargs["load_in_8bit"] = True
166                 else:
167                     raise ValueError(
168                         f"Quantization {quantization} is not supported in temporary"
169                     )
170             else:
171                 if num_gpus != 1:
172                     raise ValueError(f"Quantization is not supported for multi-gpu")
173                 elif quantization != "8-bit":
174                     raise ValueError(
175                         f"Only 8-bit quantization is supported if it is not linux system or cuda device"
176                     )
177                 else:
178                     self._model, self._tokenizer = load_compress_model(
179                         model_path=self.model_path,
180                         device=device,
181                         torch_dtype=kwargs["torch_dtype"],
182                         use_fast=self._use_fast_tokenizer,
183                         revision=kwargs["revision"],
184                     )
185                     logger.debug(f"Model Memory: {self._model.get_memory_footprint()}")
186                     return
187 
188         self._model, self._tokenizer = self._load_model(kwargs)
189 
190         if (
191             device == "cuda" and num_gpus == 1 and quantization == "none"
192         ) or device == "mps":
193             self._model.to(device)
194         logger.debug(f"Model Memory: {self._model.get_memory_footprint()}")
195 
196     @classmethod
197     def match(cls, llm_family: "LLMFamilyV1", llm_spec: "LLMSpecV1") -> bool:
198         if llm_spec.model_format != "pytorch":
199             return False
200         if "baichuan" in llm_family.model_name:
201             return False
202         if "generate" not in llm_family.model_ability:
203             return False
204         return True
205 
206     def generate(
207         self, prompt: str, generate_config: Optional[PytorchGenerateConfig] = None
208     ) -> Union[Completion, Iterator[CompletionChunk]]:
209         from .utils import generate_stream
210 
211         def generator_wrapper(
212             prompt: str, device: str, generate_config: PytorchGenerateConfig
213         ) -> Iterator[CompletionChunk]:
214             for completion_chunk, _ in generate_stream(
215                 self._model, self._tokenizer, prompt, device, generate_config
216             ):
217                 yield completion_chunk
218 
219         logger.debug(
220             "Enter generate, prompt: %s, generate config: %s", prompt, generate_config
221         )
222 
223         generate_config = self._sanitize_generate_config(generate_config)
224 
225         assert self._model is not None
226         assert self._tokenizer is not None
227 
228         stream = generate_config.get("stream", False)
229         if self._is_darwin_and_apple_silicon():
230             device = self._pytorch_model_config.get("device", "mps")
231         else:
232             device = self._pytorch_model_config.get("device", "cuda")
233         if not stream:
234             for completion_chunk, completion_usage in generate_stream(
235                 self._model, self._tokenizer, prompt, device, generate_config
236             ):
237                 pass
238             completion = Completion(
239                 id=completion_chunk["id"],
240                 object=completion_chunk["object"],
241                 created=completion_chunk["created"],
242                 model=completion_chunk["model"],
243                 choices=completion_chunk["choices"],
244                 usage=completion_usage,
245             )
246             return completion
247         else:
248             return generator_wrapper(prompt, device, generate_config)
249 
250     def create_embedding(self, input: Union[str, List[str]]) -> Embedding:
251         raise NotImplementedError
252 
253 
254 class PytorchChatModel(PytorchModel, ChatModelMixin):
255     def __init__(
256         self,
257         model_uid: str,
258         model_family: "LLMFamilyV1",
259         model_spec: "LLMSpecV1",
260         quantization: str,
261         model_path: str,
262         pytorch_model_config: Optional[PytorchModelConfig] = None,
263     ):
264         super().__init__(
265             model_uid,
266             model_family,
267             model_spec,
268             quantization,
269             model_path,
270             pytorch_model_config,
271         )
272 
273     def _sanitize_generate_config(
274         self,
275         pytorch_generate_config: Optional[PytorchGenerateConfig],
276     ) -> PytorchGenerateConfig:
277         pytorch_generate_config = super()._sanitize_generate_config(
278             pytorch_generate_config
279         )
280         if (
281             "stop" not in pytorch_generate_config
282             and self.model_family.prompt_style
283             and self.model_family.prompt_style.stop
284         ):
285             pytorch_generate_config["stop"] = self.model_family.prompt_style.stop
286         if (
287             "stop_token_ids" not in pytorch_generate_config
288             and self.model_family.prompt_style
289             and self.model_family.prompt_style.stop_token_ids
290         ):
291             pytorch_generate_config[
292                 "stop_token_ids"
293             ] = self.model_family.prompt_style.stop_token_ids
294 
295         return pytorch_generate_config
296 
297     @classmethod
298     def match(cls, llm_family: "LLMFamilyV1", llm_spec: "LLMSpecV1") -> bool:
299         if llm_spec.model_format != "pytorch":
300             return False
301         if "baichuan" in llm_family.model_name:
302             return False
303         if "chat" not in llm_family.model_ability:
304             return False
305         return True
306 
307     def chat(
308         self,
309         prompt: str,
310         system_prompt: Optional[str] = None,
311         chat_history: Optional[List[ChatCompletionMessage]] = None,
312         generate_config: Optional[PytorchGenerateConfig] = None,
313     ) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:
314         assert self.model_family.prompt_style is not None
315         prompt_style = self.model_family.prompt_style.copy()
316         if system_prompt:
317             prompt_style.system_prompt = system_prompt
318         chat_history = chat_history or []
319         full_prompt = self.get_prompt(prompt, chat_history, prompt_style)
320 
321         generate_config = self._sanitize_generate_config(generate_config)
322 
323         stream = generate_config.get("stream", False)
324         if stream:
325             it = self.generate(full_prompt, generate_config)
326             assert isinstance(it, Iterator)
327             return self._convert_chat_completion_chunks_to_chat(it)
328         else:
329             c = self.generate(full_prompt, generate_config)
330             assert not isinstance(c, Iterator)
331             return self._convert_text_completion_to_chat(c)
332
```
--- 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/pytorch/core.py b/xinference/model/llm/pytorch/core.py
--- a/xinference/model/llm/pytorch/core.py
+++ b/xinference/model/llm/pytorch/core.py
@@ -23,6 +23,8 @@
     Completion,
     CompletionChunk,
     Embedding,
+    EmbeddingData,
+    EmbeddingUsage,
 )
 from ..core import LLM
 from ..llm_family import LLMFamilyV1, LLMSpecV1
@@ -248,7 +250,85 @@
             return generator_wrapper(prompt, device, generate_config)
 
     def create_embedding(self, input: Union[str, List[str]]) -> Embedding:
-        raise NotImplementedError
+        try:
+            import torch
+            import torch.nn.functional as F
+        except ImportError as e:
+            raise ImportError(
+                "Could not import torch. Please install it with `pip install torch`."
+            ) from e
+
+        if self._is_darwin_and_apple_silicon():
+            device = self._pytorch_model_config.get("device", "mps")
+        else:
+            device = self._pytorch_model_config.get("device", "cuda")
+
+        if isinstance(input, str):
+            inputs = [input]
+        else:
+            inputs = input
+
+        tokenizer = self._tokenizer
+        is_llama = "llama" in str(type(self._model))  # llama supports batch inference
+        is_chatglm = "chatglm" in str(type(self._model))
+        if is_llama:
+            encoding = tokenizer.batch_encode_plus(
+                inputs, padding=True, return_tensors="pt"
+            )
+            input_ids = encoding["input_ids"].to(device)
+            attention_mask = encoding["attention_mask"].to(device)
+            model_output = self._model(
+                input_ids, attention_mask, output_hidden_states=True
+            )
+            data = model_output.hidden_states[-1]
+            mask = attention_mask.unsqueeze(-1).expand(data.size()).float()
+            masked_embeddings = data * mask
+            sum_embeddings = torch.sum(masked_embeddings, dim=1)
+            seq_length = torch.sum(mask, dim=1)
+            embedding = sum_embeddings / seq_length
+            normalized_embeddings = F.normalize(embedding, p=2, dim=1)
+            normalized_embeddings = normalized_embeddings.tolist()
+            token_num = torch.sum(attention_mask).item()
+
+            embedding_list = []
+            for index, data in enumerate(normalized_embeddings):
+                embedding_list.append(
+                    EmbeddingData(index=index, object="embedding", embedding=data)
+                )
+
+            usage = EmbeddingUsage(prompt_tokens=token_num, total_tokens=token_num)
+
+            ret = Embedding(
+                object="list",
+                model=self.model_uid,
+                data=embedding_list,
+                usage=usage,
+            )
+
+        else:
+            embedding = []
+            token_num = 0
+            for index, text in enumerate(inputs):
+                input_ids = tokenizer.encode(text, return_tensors="pt").to(device)
+                model_output = self._model(input_ids, output_hidden_states=True)
+                if is_chatglm:
+                    data = (model_output.hidden_states[-1].transpose(0, 1))[0]
+                else:
+                    data = model_output.hidden_states[-1][0]
+                data = F.normalize(torch.mean(data, dim=0), p=2, dim=0)
+                data = data.tolist()
+
+                embedding.append(
+                    EmbeddingData(index=index, object="embedding", embedding=data)
+                )
+                token_num += len(input_ids[0])
+
+            usage = EmbeddingUsage(prompt_tokens=token_num, total_tokens=token_num)
+            ret = Embedding(
+                object="list", model=self.model_uid, data=embedding, usage=usage
+            )
+
+        return ret
 
 
 class PytorchChatModel(PytorchModel, ChatModelMixin):

{"golden_diff": "diff --git a/xinference/model/llm/pytorch/core.py b/xinference/model/llm/pytorch/core.py\n--- a/xinference/model/llm/pytorch/core.py\n+++ b/xinference/model/llm/pytorch/core.py\n@@ -23,6 +23,8 @@\n     Completion,\n     CompletionChunk,\n     Embedding,\n+    EmbeddingData,\n+    EmbeddingUsage,\n )\n from ..core import LLM\n from ..llm_family import LLMFamilyV1, LLMSpecV1\n@@ -248,7 +250,85 @@\n             return generator_wrapper(prompt, device, generate_config)\n \n     def create_embedding(self, input: Union[str, List[str]]) -> Embedding:\n-        raise NotImplementedError\n+        try:\n+            import torch\n+            import torch.nn.functional as F\n+        except ImportError as e:\n+            raise ImportError(\n+                \"Could not import torch. Please install it with `pip install torch`.\"\n+            ) from e\n+\n+        if self._is_darwin_and_apple_silicon():\n+            device = self._pytorch_model_config.get(\"device\", \"mps\")\n+        else:\n+            device = self._pytorch_model_config.get(\"device\", \"cuda\")\n+\n+        if isinstance(input, str):\n+            inputs = [input]\n+        else:\n+            inputs = input\n+\n+        tokenizer = self._tokenizer\n+        is_llama = \"llama\" in str(type(self._model))  # llama supports batch inference\n+        is_chatglm = \"chatglm\" in str(type(self._model))\n+        if is_llama:\n+            encoding = tokenizer.batch_encode_plus(\n+                inputs, padding=True, return_tensors=\"pt\"\n+            )\n+            input_ids = encoding[\"input_ids\"].to(device)\n+            attention_mask = encoding[\"attention_mask\"].to(device)\n+            model_output = self._model(\n+                input_ids, attention_mask, output_hidden_states=True\n+            )\n+            data = model_output.hidden_states[-1]\n+            mask = attention_mask.unsqueeze(-1).expand(data.size()).float()\n+            masked_embeddings = data * mask\n+            sum_embeddings = torch.sum(masked_embeddings, dim=1)\n+            seq_length = torch.sum(mask, dim=1)\n+            embedding = sum_embeddings / seq_length\n+            normalized_embeddings = F.normalize(embedding, p=2, dim=1)\n+            normalized_embeddings = normalized_embeddings.tolist()\n+            token_num = torch.sum(attention_mask).item()\n+\n+            embedding_list = []\n+            for index, data in enumerate(normalized_embeddings):\n+                embedding_list.append(\n+                    EmbeddingData(index=index, object=\"embedding\", embedding=data)\n+                )\n+\n+            usage = EmbeddingUsage(prompt_tokens=token_num, total_tokens=token_num)\n+\n+            ret = Embedding(\n+                object=\"list\",\n+                model=self.model_uid,\n+                data=embedding_list,\n+                usage=usage,\n+            )\n+\n+        else:\n+            embedding = []\n+            token_num = 0\n+            for index, text in enumerate(inputs):\n+                input_ids = tokenizer.encode(text, return_tensors=\"pt\").to(device)\n+                model_output = self._model(input_ids, output_hidden_states=True)\n+                if is_chatglm:\n+                    data = (model_output.hidden_states[-1].transpose(0, 1))[0]\n+                else:\n+                    data = model_output.hidden_states[-1][0]\n+                data = F.normalize(torch.mean(data, dim=0), p=2, dim=0)\n+                data = data.tolist()\n+\n+                embedding.append(\n+                    EmbeddingData(index=index, object=\"embedding\", embedding=data)\n+                )\n+                token_num += len(input_ids[0])\n+\n+            usage = EmbeddingUsage(prompt_tokens=token_num, total_tokens=token_num)\n+            ret = Embedding(\n+                object=\"list\", model=self.model_uid, data=embedding, usage=usage\n+            )\n+\n+        return ret\n \n \n class PytorchChatModel(PytorchModel, ChatModelMixin):\n", "issue": "FEAT: PyTorch model embeddings\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 logging\nfrom typing import Iterator, List, Optional, TypedDict, Union\n\nfrom ....constants import XINFERENCE_CACHE_DIR\nfrom ....types import (\n    ChatCompletion,\n    ChatCompletionChunk,\n    ChatCompletionMessage,\n    Completion,\n    CompletionChunk,\n    Embedding,\n)\nfrom ..core import LLM\nfrom ..llm_family import LLMFamilyV1, LLMSpecV1\nfrom ..utils import ChatModelMixin\n\nlogger = logging.getLogger(__name__)\n\n\nclass PytorchGenerateConfig(TypedDict, total=False):\n    temperature: float\n    repetition_penalty: float\n    top_p: float\n    top_k: int\n    stream: bool\n    max_new_tokens: int\n    echo: bool\n    stop: Optional[Union[str, List[str]]]\n    stop_token_ids: Optional[Union[int, List[int]]]\n    stream_interval: int\n    model: Optional[str]\n\n\nclass PytorchModelConfig(TypedDict, total=False):\n    revision: str\n    device: str\n    gpus: Optional[str]\n    num_gpus: int\n    max_gpu_memory: str\n    gptq_ckpt: Optional[str]\n    gptq_wbits: int\n    gptq_groupsize: int\n    gptq_act_order: bool\n\n\nclass PytorchModel(LLM):\n    def __init__(\n        self,\n        model_uid: str,\n        model_family: \"LLMFamilyV1\",\n        model_spec: \"LLMSpecV1\",\n        quantization: str,\n        model_path: str,\n        pytorch_model_config: Optional[PytorchModelConfig] = None,\n    ):\n        super().__init__(model_uid, model_family, model_spec, quantization, model_path)\n        self._use_fast_tokenizer = True\n        self._pytorch_model_config: PytorchModelConfig = self._sanitize_model_config(\n            pytorch_model_config\n        )\n\n    def _sanitize_model_config(\n        self, pytorch_model_config: Optional[PytorchModelConfig]\n    ) -> PytorchModelConfig:\n        if pytorch_model_config is None:\n            pytorch_model_config = PytorchModelConfig()\n        pytorch_model_config.setdefault(\"revision\", \"main\")\n        pytorch_model_config.setdefault(\"gpus\", None)\n        pytorch_model_config.setdefault(\"num_gpus\", 1)\n        pytorch_model_config.setdefault(\"gptq_ckpt\", None)\n        pytorch_model_config.setdefault(\"gptq_wbits\", 16)\n        pytorch_model_config.setdefault(\"gptq_groupsize\", -1)\n        pytorch_model_config.setdefault(\"gptq_act_order\", False)\n        if self._is_darwin_and_apple_silicon():\n            pytorch_model_config.setdefault(\"device\", \"mps\")\n        else:\n            pytorch_model_config.setdefault(\"device\", \"cuda\")\n        return pytorch_model_config\n\n    def _sanitize_generate_config(\n        self,\n        pytorch_generate_config: Optional[PytorchGenerateConfig],\n    ) -> PytorchGenerateConfig:\n        if pytorch_generate_config is None:\n            pytorch_generate_config = PytorchGenerateConfig()\n        pytorch_generate_config.setdefault(\"temperature\", 0.7)\n        pytorch_generate_config.setdefault(\"repetition_penalty\", 1.0)\n        pytorch_generate_config.setdefault(\"max_new_tokens\", 512)\n        pytorch_generate_config.setdefault(\"stream_interval\", 2)\n        pytorch_generate_config[\"model\"] = self.model_uid\n        return pytorch_generate_config\n\n    def _load_model(self, kwargs: dict):\n        try:\n            from transformers import AutoModelForCausalLM, AutoTokenizer\n        except ImportError:\n            error_message = \"Failed to import module 'transformers'\"\n            installation_guide = [\n                \"Please make sure 'transformers' is installed. \",\n                \"You can install it by `pip install transformers`\\n\",\n            ]\n\n            raise ImportError(f\"{error_message}\\n\\n{''.join(installation_guide)}\")\n\n        tokenizer = AutoTokenizer.from_pretrained(\n            self.model_path,\n            use_fast=self._use_fast_tokenizer,\n            revision=kwargs[\"revision\"],\n            cache_dir=XINFERENCE_CACHE_DIR,\n        )\n        model = AutoModelForCausalLM.from_pretrained(\n            self.model_path,\n            low_cpu_mem_usage=True,\n            cache_dir=XINFERENCE_CACHE_DIR,\n            **kwargs,\n        )\n        return model, tokenizer\n\n    def load(self):\n        try:\n            import torch\n        except ImportError:\n            raise ImportError(\n                f\"Failed to import module 'torch'. Please make sure 'torch' is installed.\\n\\n\"\n            )\n        from .compression import load_compress_model\n\n        quantization = self.quantization\n        num_gpus = self._pytorch_model_config.get(\"num_gpus\", 1)\n        if self._is_darwin_and_apple_silicon():\n            device = self._pytorch_model_config.get(\"device\", \"mps\")\n        else:\n            device = self._pytorch_model_config.get(\"device\", \"cuda\")\n\n        if device == \"cpu\":\n            kwargs = {\"torch_dtype\": torch.float32}\n        elif device == \"cuda\":\n            kwargs = {\"torch_dtype\": torch.float16}\n        elif device == \"mps\":\n            kwargs = {\"torch_dtype\": torch.float16}\n        else:\n            raise ValueError(f\"Device {device} is not supported in temporary\")\n        kwargs[\"revision\"] = self._pytorch_model_config.get(\"revision\", \"main\")\n\n        if quantization != \"none\":\n            if device == \"cuda\" and self._is_linux():\n                kwargs[\"device_map\"] = \"auto\"\n                if quantization == \"4-bit\":\n                    kwargs[\"load_in_4bit\"] = True\n                elif quantization == \"8-bit\":\n                    kwargs[\"load_in_8bit\"] = True\n                else:\n                    raise ValueError(\n                        f\"Quantization {quantization} is not supported in temporary\"\n                    )\n            else:\n                if num_gpus != 1:\n                    raise ValueError(f\"Quantization is not supported for multi-gpu\")\n                elif quantization != \"8-bit\":\n                    raise ValueError(\n                        f\"Only 8-bit quantization is supported if it is not linux system or cuda device\"\n                    )\n                else:\n                    self._model, self._tokenizer = load_compress_model(\n                        model_path=self.model_path,\n                        device=device,\n                        torch_dtype=kwargs[\"torch_dtype\"],\n                        use_fast=self._use_fast_tokenizer,\n                        revision=kwargs[\"revision\"],\n                    )\n                    logger.debug(f\"Model Memory: {self._model.get_memory_footprint()}\")\n                    return\n\n        self._model, self._tokenizer = self._load_model(kwargs)\n\n        if (\n            device == \"cuda\" and num_gpus == 1 and quantization == \"none\"\n        ) or device == \"mps\":\n            self._model.to(device)\n        logger.debug(f\"Model Memory: {self._model.get_memory_footprint()}\")\n\n    @classmethod\n    def match(cls, llm_family: \"LLMFamilyV1\", llm_spec: \"LLMSpecV1\") -> bool:\n        if llm_spec.model_format != \"pytorch\":\n            return False\n        if \"baichuan\" in llm_family.model_name:\n            return False\n        if \"generate\" not in llm_family.model_ability:\n            return False\n        return True\n\n    def generate(\n        self, prompt: str, generate_config: Optional[PytorchGenerateConfig] = None\n    ) -> Union[Completion, Iterator[CompletionChunk]]:\n        from .utils import generate_stream\n\n        def generator_wrapper(\n            prompt: str, device: str, generate_config: PytorchGenerateConfig\n        ) -> Iterator[CompletionChunk]:\n            for completion_chunk, _ in generate_stream(\n                self._model, self._tokenizer, prompt, device, generate_config\n            ):\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        assert self._model is not None\n        assert self._tokenizer is not None\n\n        stream = generate_config.get(\"stream\", False)\n        if self._is_darwin_and_apple_silicon():\n            device = self._pytorch_model_config.get(\"device\", \"mps\")\n        else:\n            device = self._pytorch_model_config.get(\"device\", \"cuda\")\n        if not stream:\n            for completion_chunk, completion_usage in generate_stream(\n                self._model, self._tokenizer, prompt, device, generate_config\n            ):\n                pass\n            completion = Completion(\n                id=completion_chunk[\"id\"],\n                object=completion_chunk[\"object\"],\n                created=completion_chunk[\"created\"],\n                model=completion_chunk[\"model\"],\n                choices=completion_chunk[\"choices\"],\n                usage=completion_usage,\n            )\n            return completion\n        else:\n            return generator_wrapper(prompt, device, generate_config)\n\n    def create_embedding(self, input: Union[str, List[str]]) -> Embedding:\n        raise NotImplementedError\n\n\nclass PytorchChatModel(PytorchModel, ChatModelMixin):\n    def __init__(\n        self,\n        model_uid: str,\n        model_family: \"LLMFamilyV1\",\n        model_spec: \"LLMSpecV1\",\n        quantization: str,\n        model_path: str,\n        pytorch_model_config: Optional[PytorchModelConfig] = None,\n    ):\n        super().__init__(\n            model_uid,\n            model_family,\n            model_spec,\n            quantization,\n            model_path,\n            pytorch_model_config,\n        )\n\n    def _sanitize_generate_config(\n        self,\n        pytorch_generate_config: Optional[PytorchGenerateConfig],\n    ) -> PytorchGenerateConfig:\n        pytorch_generate_config = super()._sanitize_generate_config(\n            pytorch_generate_config\n        )\n        if (\n            \"stop\" not in pytorch_generate_config\n            and self.model_family.prompt_style\n            and self.model_family.prompt_style.stop\n        ):\n            pytorch_generate_config[\"stop\"] = self.model_family.prompt_style.stop\n        if (\n            \"stop_token_ids\" not in pytorch_generate_config\n            and self.model_family.prompt_style\n            and self.model_family.prompt_style.stop_token_ids\n        ):\n            pytorch_generate_config[\n                \"stop_token_ids\"\n            ] = self.model_family.prompt_style.stop_token_ids\n\n        return pytorch_generate_config\n\n    @classmethod\n    def match(cls, llm_family: \"LLMFamilyV1\", llm_spec: \"LLMSpecV1\") -> bool:\n        if llm_spec.model_format != \"pytorch\":\n            return False\n        if \"baichuan\" in llm_family.model_name:\n            return False\n        if \"chat\" not in llm_family.model_ability:\n            return False\n        return True\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[PytorchGenerateConfig] = None,\n    ) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:\n        assert self.model_family.prompt_style is not None\n        prompt_style = self.model_family.prompt_style.copy()\n        if system_prompt:\n            prompt_style.system_prompt = system_prompt\n        chat_history = chat_history or []\n        full_prompt = self.get_prompt(prompt, chat_history, prompt_style)\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/pytorch/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 logging\nfrom typing import Iterator, List, Optional, TypedDict, Union\n\nfrom ....constants import XINFERENCE_CACHE_DIR\nfrom ....types import (\n    ChatCompletion,\n    ChatCompletionChunk,\n    ChatCompletionMessage,\n    Completion,\n    CompletionChunk,\n    Embedding,\n    EmbeddingData,\n    EmbeddingUsage,\n)\nfrom ..core import LLM\nfrom ..llm_family import LLMFamilyV1, LLMSpecV1\nfrom ..utils import ChatModelMixin\n\nlogger = logging.getLogger(__name__)\n\n\nclass PytorchGenerateConfig(TypedDict, total=False):\n    temperature: float\n    repetition_penalty: float\n    top_p: float\n    top_k: int\n    stream: bool\n    max_new_tokens: int\n    echo: bool\n    stop: Optional[Union[str, List[str]]]\n    stop_token_ids: Optional[Union[int, List[int]]]\n    stream_interval: int\n    model: Optional[str]\n\n\nclass PytorchModelConfig(TypedDict, total=False):\n    revision: str\n    device: str\n    gpus: Optional[str]\n    num_gpus: int\n    max_gpu_memory: str\n    gptq_ckpt: Optional[str]\n    gptq_wbits: int\n    gptq_groupsize: int\n    gptq_act_order: bool\n\n\nclass PytorchModel(LLM):\n    def __init__(\n        self,\n        model_uid: str,\n        model_family: \"LLMFamilyV1\",\n        model_spec: \"LLMSpecV1\",\n        quantization: str,\n        model_path: str,\n        pytorch_model_config: Optional[PytorchModelConfig] = None,\n    ):\n        super().__init__(model_uid, model_family, model_spec, quantization, model_path)\n        self._use_fast_tokenizer = True\n        self._pytorch_model_config: PytorchModelConfig = self._sanitize_model_config(\n            pytorch_model_config\n        )\n\n    def _sanitize_model_config(\n        self, pytorch_model_config: Optional[PytorchModelConfig]\n    ) -> PytorchModelConfig:\n        if pytorch_model_config is None:\n            pytorch_model_config = PytorchModelConfig()\n        pytorch_model_config.setdefault(\"revision\", \"main\")\n        pytorch_model_config.setdefault(\"gpus\", None)\n        pytorch_model_config.setdefault(\"num_gpus\", 1)\n        pytorch_model_config.setdefault(\"gptq_ckpt\", None)\n        pytorch_model_config.setdefault(\"gptq_wbits\", 16)\n        pytorch_model_config.setdefault(\"gptq_groupsize\", -1)\n        pytorch_model_config.setdefault(\"gptq_act_order\", False)\n        if self._is_darwin_and_apple_silicon():\n            pytorch_model_config.setdefault(\"device\", \"mps\")\n        else:\n            pytorch_model_config.setdefault(\"device\", \"cuda\")\n        return pytorch_model_config\n\n    def _sanitize_generate_config(\n        self,\n        pytorch_generate_config: Optional[PytorchGenerateConfig],\n    ) -> PytorchGenerateConfig:\n        if pytorch_generate_config is None:\n            pytorch_generate_config = PytorchGenerateConfig()\n        pytorch_generate_config.setdefault(\"temperature\", 0.7)\n        pytorch_generate_config.setdefault(\"repetition_penalty\", 1.0)\n        pytorch_generate_config.setdefault(\"max_new_tokens\", 512)\n        pytorch_generate_config.setdefault(\"stream_interval\", 2)\n        pytorch_generate_config[\"model\"] = self.model_uid\n        return pytorch_generate_config\n\n    def _load_model(self, kwargs: dict):\n        try:\n            from transformers import AutoModelForCausalLM, AutoTokenizer\n        except ImportError:\n            error_message = \"Failed to import module 'transformers'\"\n            installation_guide = [\n                \"Please make sure 'transformers' is installed. \",\n                \"You can install it by `pip install transformers`\\n\",\n            ]\n\n            raise ImportError(f\"{error_message}\\n\\n{''.join(installation_guide)}\")\n\n        tokenizer = AutoTokenizer.from_pretrained(\n            self.model_path,\n            use_fast=self._use_fast_tokenizer,\n            revision=kwargs[\"revision\"],\n            cache_dir=XINFERENCE_CACHE_DIR,\n        )\n        model = AutoModelForCausalLM.from_pretrained(\n            self.model_path,\n            low_cpu_mem_usage=True,\n            cache_dir=XINFERENCE_CACHE_DIR,\n            **kwargs,\n        )\n        return model, tokenizer\n\n    def load(self):\n        try:\n            import torch\n        except ImportError:\n            raise ImportError(\n                f\"Failed to import module 'torch'. Please make sure 'torch' is installed.\\n\\n\"\n            )\n        from .compression import load_compress_model\n\n        quantization = self.quantization\n        num_gpus = self._pytorch_model_config.get(\"num_gpus\", 1)\n        if self._is_darwin_and_apple_silicon():\n            device = self._pytorch_model_config.get(\"device\", \"mps\")\n        else:\n            device = self._pytorch_model_config.get(\"device\", \"cuda\")\n\n        if device == \"cpu\":\n            kwargs = {\"torch_dtype\": torch.float32}\n        elif device == \"cuda\":\n            kwargs = {\"torch_dtype\": torch.float16}\n        elif device == \"mps\":\n            kwargs = {\"torch_dtype\": torch.float16}\n        else:\n            raise ValueError(f\"Device {device} is not supported in temporary\")\n        kwargs[\"revision\"] = self._pytorch_model_config.get(\"revision\", \"main\")\n\n        if quantization != \"none\":\n            if device == \"cuda\" and self._is_linux():\n                kwargs[\"device_map\"] = \"auto\"\n                if quantization == \"4-bit\":\n                    kwargs[\"load_in_4bit\"] = True\n                elif quantization == \"8-bit\":\n                    kwargs[\"load_in_8bit\"] = True\n                else:\n                    raise ValueError(\n                        f\"Quantization {quantization} is not supported in temporary\"\n                    )\n            else:\n                if num_gpus != 1:\n                    raise ValueError(f\"Quantization is not supported for multi-gpu\")\n                elif quantization != \"8-bit\":\n                    raise ValueError(\n                        f\"Only 8-bit quantization is supported if it is not linux system or cuda device\"\n                    )\n                else:\n                    self._model, self._tokenizer = load_compress_model(\n                        model_path=self.model_path,\n                        device=device,\n                        torch_dtype=kwargs[\"torch_dtype\"],\n                        use_fast=self._use_fast_tokenizer,\n                        revision=kwargs[\"revision\"],\n                    )\n                    logger.debug(f\"Model Memory: {self._model.get_memory_footprint()}\")\n                    return\n\n        self._model, self._tokenizer = self._load_model(kwargs)\n\n        if (\n            device == \"cuda\" and num_gpus == 1 and quantization == \"none\"\n        ) or device == \"mps\":\n            self._model.to(device)\n        logger.debug(f\"Model Memory: {self._model.get_memory_footprint()}\")\n\n    @classmethod\n    def match(cls, llm_family: \"LLMFamilyV1\", llm_spec: \"LLMSpecV1\") -> bool:\n        if llm_spec.model_format != \"pytorch\":\n            return False\n        if \"baichuan\" in llm_family.model_name:\n            return False\n        if \"generate\" not in llm_family.model_ability:\n            return False\n        return True\n\n    def generate(\n        self, prompt: str, generate_config: Optional[PytorchGenerateConfig] = None\n    ) -> Union[Completion, Iterator[CompletionChunk]]:\n        from .utils import generate_stream\n\n        def generator_wrapper(\n            prompt: str, device: str, generate_config: PytorchGenerateConfig\n        ) -> Iterator[CompletionChunk]:\n            for completion_chunk, _ in generate_stream(\n                self._model, self._tokenizer, prompt, device, generate_config\n            ):\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        assert self._model is not None\n        assert self._tokenizer is not None\n\n        stream = generate_config.get(\"stream\", False)\n        if self._is_darwin_and_apple_silicon():\n            device = self._pytorch_model_config.get(\"device\", \"mps\")\n        else:\n            device = self._pytorch_model_config.get(\"device\", \"cuda\")\n        if not stream:\n            for completion_chunk, completion_usage in generate_stream(\n                self._model, self._tokenizer, prompt, device, generate_config\n            ):\n                pass\n            completion = Completion(\n                id=completion_chunk[\"id\"],\n                object=completion_chunk[\"object\"],\n                created=completion_chunk[\"created\"],\n                model=completion_chunk[\"model\"],\n                choices=completion_chunk[\"choices\"],\n                usage=completion_usage,\n            )\n            return completion\n        else:\n            return generator_wrapper(prompt, device, generate_config)\n\n    def create_embedding(self, input: Union[str, List[str]]) -> Embedding:\n        try:\n            import torch\n            import torch.nn.functional as F\n        except ImportError as e:\n            raise ImportError(\n                \"Could not import torch. Please install it with `pip install torch`.\"\n            ) from e\n\n        if self._is_darwin_and_apple_silicon():\n            device = self._pytorch_model_config.get(\"device\", \"mps\")\n        else:\n            device = self._pytorch_model_config.get(\"device\", \"cuda\")\n\n        if isinstance(input, str):\n            inputs = [input]\n        else:\n            inputs = input\n\n        tokenizer = self._tokenizer\n        is_llama = \"llama\" in str(type(self._model))  # llama supports batch inference\n        is_chatglm = \"chatglm\" in str(type(self._model))\n        if is_llama:\n            encoding = tokenizer.batch_encode_plus(\n                inputs, padding=True, return_tensors=\"pt\"\n            )\n            input_ids = encoding[\"input_ids\"].to(device)\n            attention_mask = encoding[\"attention_mask\"].to(device)\n            model_output = self._model(\n                input_ids, attention_mask, output_hidden_states=True\n            )\n            data = model_output.hidden_states[-1]\n            mask = attention_mask.unsqueeze(-1).expand(data.size()).float()\n            masked_embeddings = data * mask\n            sum_embeddings = torch.sum(masked_embeddings, dim=1)\n            seq_length = torch.sum(mask, dim=1)\n            embedding = sum_embeddings / seq_length\n            normalized_embeddings = F.normalize(embedding, p=2, dim=1)\n            normalized_embeddings = normalized_embeddings.tolist()\n            token_num = torch.sum(attention_mask).item()\n\n            embedding_list = []\n            for index, data in enumerate(normalized_embeddings):\n                embedding_list.append(\n                    EmbeddingData(index=index, object=\"embedding\", embedding=data)\n                )\n\n            usage = EmbeddingUsage(prompt_tokens=token_num, total_tokens=token_num)\n\n            ret = Embedding(\n                object=\"list\",\n                model=self.model_uid,\n                data=embedding_list,\n                usage=usage,\n            )\n\n        else:\n            embedding = []\n            token_num = 0\n            for index, text in enumerate(inputs):\n                input_ids = tokenizer.encode(text, return_tensors=\"pt\").to(device)\n                model_output = self._model(input_ids, output_hidden_states=True)\n                if is_chatglm:\n                    data = (model_output.hidden_states[-1].transpose(0, 1))[0]\n                else:\n                    data = model_output.hidden_states[-1][0]\n                data = F.normalize(torch.mean(data, dim=0), p=2, dim=0)\n                data = data.tolist()\n\n                embedding.append(\n                    EmbeddingData(index=index, object=\"embedding\", embedding=data)\n                )\n                token_num += len(input_ids[0])\n\n            usage = EmbeddingUsage(prompt_tokens=token_num, total_tokens=token_num)\n            ret = Embedding(\n                object=\"list\", model=self.model_uid, data=embedding, usage=usage\n            )\n\n        return ret\n\n\nclass PytorchChatModel(PytorchModel, ChatModelMixin):\n    def __init__(\n        self,\n        model_uid: str,\n        model_family: \"LLMFamilyV1\",\n        model_spec: \"LLMSpecV1\",\n        quantization: str,\n        model_path: str,\n        pytorch_model_config: Optional[PytorchModelConfig] = None,\n    ):\n        super().__init__(\n            model_uid,\n            model_family,\n            model_spec,\n            quantization,\n            model_path,\n            pytorch_model_config,\n        )\n\n    def _sanitize_generate_config(\n        self,\n        pytorch_generate_config: Optional[PytorchGenerateConfig],\n    ) -> PytorchGenerateConfig:\n        pytorch_generate_config = super()._sanitize_generate_config(\n            pytorch_generate_config\n        )\n        if (\n            \"stop\" not in pytorch_generate_config\n            and self.model_family.prompt_style\n            and self.model_family.prompt_style.stop\n        ):\n            pytorch_generate_config[\"stop\"] = self.model_family.prompt_style.stop\n        if (\n            \"stop_token_ids\" not in pytorch_generate_config\n            and self.model_family.prompt_style\n            and self.model_family.prompt_style.stop_token_ids\n        ):\n            pytorch_generate_config[\n                \"stop_token_ids\"\n            ] = self.model_family.prompt_style.stop_token_ids\n\n        return pytorch_generate_config\n\n    @classmethod\n    def match(cls, llm_family: \"LLMFamilyV1\", llm_spec: \"LLMSpecV1\") -> bool:\n        if llm_spec.model_format != \"pytorch\":\n            return False\n        if \"baichuan\" in llm_family.model_name:\n            return False\n        if \"chat\" not in llm_family.model_ability:\n            return False\n        return True\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[PytorchGenerateConfig] = None,\n    ) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:\n        assert self.model_family.prompt_style is not None\n        prompt_style = self.model_family.prompt_style.copy()\n        if system_prompt:\n            prompt_style.system_prompt = system_prompt\n        chat_history = chat_history or []\n        full_prompt = self.get_prompt(prompt, chat_history, prompt_style)\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/pytorch/core.py"}]}