This commit is contained in:
Andrei Betlen 2023-04-03 18:46:19 -04:00
parent 7d1977e8f0
commit ae004eb69e

View file

@ -262,6 +262,17 @@ class Llama:
}, },
} }
def embed(self, input: str) -> List[float]:
"""Embed a string.
Args:
input: The utf-8 encoded string to embed.
Returns:
A list of embeddings
"""
return list(map(float, self.create_embedding(input)["data"][0]["embedding"]))
def _create_completion( def _create_completion(
self, self,
prompt: str, prompt: str,
@ -341,7 +352,7 @@ class Llama:
"model": self.model_path, "model": self.model_path,
"choices": [ "choices": [
{ {
"text": text[start :].decode("utf-8"), "text": text[start:].decode("utf-8"),
"index": 0, "index": 0,
"logprobs": None, "logprobs": None,
"finish_reason": None, "finish_reason": None,