Embeddings api (#59)

* added Embeddings class to embeddings.py for embeddings API

* Added Embeddings API to python library and CLI, with documentation for both in the README.md

also changed the version number

* includes black ruff and mypy fixes

* Changed default embeddings model to bert base uncased

---------

Co-authored-by: Carson <[email protected]>

Author: clam004

README.md

@@ -453,6 +453,16 @@ print(output_text)
 Space Robots are a great way to get your kids interested in science. After all, they are the future!
 ```
 
+## Embeddings API
+
+Embeddings are vector representations of sequences. You can use these vectors for measuring the overall similarity between texts. Embeddings are useful for tasks such as search and retrieval. 
+
+```python
+resp = together.Embeddings.create("embed this sentence into a single vector", model="togethercomputer/bert-base-uncased")
+
+print(resp['data'][0]['embedding']) # [0.06659205, 0.07896972, 0.007910785 ........]
+```
+
 ## Colab Tutorial
 
 Follow along in our Colab (Google Colaboratory) Notebook Tutorial [Example Finetuning Project](https://colab.research.google.com/drive/11DwtftycpDSgp3Z1vnV-Cy68zvkGZL4K?usp=sharing).
@@ -512,6 +522,9 @@ together finetune download ft-dd93c727-f35e-41c2-a370-7d55b54128fa
 
 # inference using your new finetuned model (with new finetuned model name from together models list)
 together complete "Space robots are" -m yourname/ft-dd93c727-f35e-41c2-a370-7d55b54128fa-2023-08-16-10-15-09
+
+# create an embedding from your input sequence
+together embeddings "embed this sentence into a single vector" -m togethercomputer/bert-base-uncased
 ```
 
 ## Contributing

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.masonry.api"
 
 [tool.poetry]
 name = "together"
-version = "0.2.6"
+version = "0.2.7"
 authors = [
   "Together AI <[email protected]>"
 ]

src/together/__init__.py

@@ -16,9 +16,11 @@
 api_base_files = urllib.parse.urljoin(api_base, "/v1/files/")
 api_base_finetune = urllib.parse.urljoin(api_base, "/v1/fine-tunes/")
 api_base_instances = urllib.parse.urljoin(api_base, "instances/")
+api_base_embeddings = urllib.parse.urljoin(api_base, "api/v1/embeddings")
 
 default_text_model = "togethercomputer/RedPajama-INCITE-7B-Chat"
 default_image_model = "runwayml/stable-diffusion-v1-5"
+default_embedding_model = "togethercomputer/bert-base-uncased"
 log_level = "WARNING"
 
 MISSING_API_KEY_MESSAGE = """TOGETHER_API_KEY not found.
@@ -31,6 +33,7 @@
 min_samples = 100
 
 from .complete import Complete
+from .embeddings import Embeddings
 from .error import *
 from .files import Files
 from .finetune import Finetune
@@ -45,13 +48,16 @@
     "api_base_files",
     "api_base_finetune",
     "api_base_instances",
+    "api_base_embeddings",
     "default_text_model",
     "default_image_model",
+    "default_embedding_model",
     "Models",
     "Complete",
     "Files",
     "Finetune",
     "Image",
+    "Embeddings",
     "MAX_CONNECTION_RETRIES",
     "MISSING_API_KEY_MESSAGE",
     "BACKOFF_FACTOR",

src/together/cli/cli.py

@@ -2,7 +2,7 @@
 import argparse
 
 import together
-from together.commands import chat, complete, files, finetune, image, models
+from together.commands import chat, complete, embeddings, files, finetune, image, models
 from together.utils import get_logger
 
 
@@ -49,6 +49,7 @@ def main() -> None:
     image.add_parser(subparser)
     files.add_parser(subparser)
     finetune.add_parser(subparser)
+    embeddings.add_parser(subparser)
 
     args = parser.parse_args()
 

src/together/commands/embeddings.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import argparse
+import json
+
+import together
+from together import Embeddings
+from together.utils import get_logger
+
+
+logger = get_logger(str(__name__))
+
+
+def add_parser(subparsers: argparse._SubParsersAction[argparse.ArgumentParser]) -> None:
+    COMMAND_NAME = "embeddings"
+    subparser = subparsers.add_parser(COMMAND_NAME)
+
+    subparser.add_argument(
+        "input",
+        metavar="INPUT",
+        default=None,
+        type=str,
+        help="A string providing context for the model to embed",
+    )
+
+    subparser.add_argument(
+        "--model",
+        "-m",
+        default=together.default_embedding_model,
+        type=str,
+        help=f"The name of the model to query. Default={together.default_text_model}",
+    )
+    subparser.set_defaults(func=_run_complete)
+
+
+def _run_complete(args: argparse.Namespace) -> None:
+    embeddings = Embeddings()
+
+    try:
+        response = embeddings.create(
+            input=args.input,
+            model=args.model,
+        )
+
+        print(json.dumps(response, indent=4))
+    except together.AuthenticationError:
+        logger.critical(together.MISSING_API_KEY_MESSAGE)
+        exit(0)

src/together/embeddings.py

@@ -0,0 +1,35 @@
+from typing import Any, Dict, Optional
+
+import together
+from together.utils import create_post_request, get_logger
+
+
+logger = get_logger(str(__name__))
+
+
+class Embeddings:
+    @classmethod
+    def create(
+        self,
+        input: str,
+        model: Optional[str] = "",
+    ) -> Dict[str, Any]:
+        if model == "":
+            model = together.default_embedding_model
+
+        parameter_payload = {
+            "input": input,
+            "model": model,
+        }
+
+        # send request
+        response = create_post_request(
+            url=together.api_base_embeddings, json=parameter_payload
+        )
+
+        try:
+            response_json = dict(response.json())
+
+        except Exception as e:
+            raise together.JSONError(e, http_status=response.status_code)
+        return response_json