Embed a string. Args: input: The utf-8 encoded string to embed. Returns: A list of embeddings
(
self,
input: Union[str, List[str]],
normalize: bool = False,
truncate: bool = True,
return_count: bool = False,
)
| 1058 | } |
| 1059 | |
| 1060 | def embed( |
| 1061 | self, |
| 1062 | input: Union[str, List[str]], |
| 1063 | normalize: bool = False, |
| 1064 | truncate: bool = True, |
| 1065 | return_count: bool = False, |
| 1066 | ): |
| 1067 | """Embed a string. |
| 1068 | |
| 1069 | Args: |
| 1070 | input: The utf-8 encoded string to embed. |
| 1071 | |
| 1072 | Returns: |
| 1073 | A list of embeddings |
| 1074 | """ |
| 1075 | n_embd = self.n_embd() |
| 1076 | n_batch = self.n_batch |
| 1077 | n_seq_max = self.context_params.n_seq_max |
| 1078 | |
| 1079 | # get pooling information |
| 1080 | pooling_type = self.pooling_type() |
| 1081 | # In embedding mode every input token must be marked as an output, regardless of |
| 1082 | # pooling type. llama.cpp would otherwise override per-token `logits[i]` and emit |
| 1083 | # "embeddings required but some input tokens were not marked as outputs -> |
| 1084 | # overriding" once per input. Pooling NONE vs MEAN/CLS only changes how the |
| 1085 | # per-token outputs are read back (see decode_batch below), not whether they are |
| 1086 | # produced. See abetlen/llama-cpp-python#2208. |
| 1087 | logits_all = True |
| 1088 | |
| 1089 | if self.context_params.embeddings is False: |
| 1090 | raise RuntimeError( |
| 1091 | "Llama model must be created with embedding=True to call this method" |
| 1092 | ) |
| 1093 | |
| 1094 | if self.verbose: |
| 1095 | llama_cpp.llama_perf_context_reset(self._ctx.ctx) |
| 1096 | |
| 1097 | if isinstance(input, str): |
| 1098 | inputs = [input] |
| 1099 | else: |
| 1100 | inputs = input |
| 1101 | |
| 1102 | # reset batch |
| 1103 | self._batch.reset() |
| 1104 | |
| 1105 | # decode and fetch embeddings |
| 1106 | data: Union[List[List[float]], List[List[List[float]]]] = [] |
| 1107 | |
| 1108 | def decode_batch(seq_sizes: List[int]): |
| 1109 | self._ctx.kv_cache_clear() |
| 1110 | self._ctx.decode(self._batch) |
| 1111 | self._batch.reset() |
| 1112 | |
| 1113 | # store embeddings |
| 1114 | if pooling_type == llama_cpp.LLAMA_POOLING_TYPE_NONE: |
| 1115 | pos: int = 0 |
| 1116 | for i, size in enumerate(seq_sizes): |
| 1117 | ptr = llama_cpp.llama_get_embeddings(self._ctx.ctx) |