Add support to get embeddings from high-level api. Closes abetlen#4

abetlen · abetlen · commit 70b8a1ef75ca · 2023-03-28T04:59:54.000-04:00
diff --git a/examples/high_level_api_embedding.py b/examples/high_level_api_embedding.py
@@ -0,0 +1,12 @@
+import json
+import argparse
+
+from llama_cpp import Llama
+
+parser = argparse.ArgumentParser()
+parser.add_argument("-m", "--model", type=str, default=".//models/...")
+args = parser.parse_args()
+
+llm = Llama(model_path=args.model, embedding=True)
+
+print(llm.embed("Hello world!"))
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -105,6 +105,20 @@ def detokenize(self, tokens: List[int]) -> bytes:
             output += llama_cpp.llama_token_to_str(self.ctx, token)
         return output
 
+    def embed(self, text: str):
+        """Embed a string.
+
+        Args:
+            text: The utf-8 encoded string to embed.
+
+        Returns:
+            A list of embeddings.
+        """
+        tokens = self.tokenize(text.encode("utf-8"))
+        self._eval(tokens, 0)
+        embeddings = llama_cpp.llama_get_embeddings(self.ctx)
+        return embeddings[:llama_cpp.llama_n_embd(self.ctx)]
+
     def _eval(self, tokens: List[int], n_past):
         rc = llama_cpp.llama_eval(
             self.ctx,
