Introducing the Together Embeddings endpoint — Higher accuracy, longer context, and lower cost

from openai import OpenAI
client = OpenAI(api_key=TOGETHER_API_KEY, base_url="https://api.together.xyz/v1")

def get_embedding(text, model="togethercomputer/m2-bert-80M-32k-retrieval"):
   text = text.replace("\n", " ")
   return client.embeddings.create(input = [text], model=model).data[0].embedding

import json, os

import numpy as np
import together

together.api_key = 

def generate_embeddings(input_texts: List[str], model_api_string: str) -> List[List[float]]:
    """Generate embeddings from Together API.

    Args:
        input_texts: a list of string input texts.
        model_api_string: str. An API string for a specific embedding model of your choice.

    Returns:
        embeddings_list: a list of embeddings. Each element corresponds to the each input text.
    """
    client = together.Together()
    outputs = client.embeddings.create(
        input=input_texts, 
        model=model_api_string,
    )
    return [x.embedding for x in outputs.data]

def sample_and_extract_embeddings(data_path: str, model_api_string: str, num_samples: int = -1, context_length=512) -> np.ndarray:
    """Sample data examples and extract embeddings for each example.

    Args:
        data_path: str. A path to the data file. It should be in the .jsonl format.
        model_api_string: str. An API string for a specific embedding model of your choice.
        num_samples: int. The number of data examples to sample.
        context_length: int. The max context length of the model (model_api_string).

    Returns:
        embeddings: np.ndarray with num_samples by m where m is the embedding dimension.

    """
    max_num_chars = context_length * 4 # Assuming that each token is about 4 characters.
    embeddings = []
    count = 0
    print(f"Reading from {data_path}")
    with open(data_path, 'r') as f:
        for line in f:
            try:
                ex = json.loads(line)
                total_chars = len(ex['text'])
                if total_chars < max_num_chars:
                    text_ls = [ex['text']]
                else:
                    text_ls = [ex['text'][i : i + max_num_chars] for i in range(0, total_chars, max_num_chars)]                
                embeddings.extend(generate_embeddings_batch(text_ls[:min(len(text_ls), num_samples-count)], model_api_string))
                count += min(len(text_ls), num_samples-count)
            except Exception as e:
                print(f'Error occurred while loading the JSON file of {data_path} with the error message {e}.')
            if count >= num_samples:
                break
    return np.array(embeddings)

import os
import argparse
from typing import List

import matplotlib
from matplotlib import pyplot as plt
from sklearn.manifold import TSNE
import numpy as np

def transform_tSNE(arr: np.ndarray, n_components: int, perplexity: int) -> np.ndarray:
    """Transform the given ndarray using a tSNE model.

    Args:
        arr: np.ndarray. In this example, an embedding matrix (n by m), where n is the number of examples and m equals to the embedding dimension.
        n_components: int. The number of components for tSNE.
        perplexity: int. Perplexity for tSNE.

    Returns:
        vis_dims: np.ndarray. A transformed matrix of n by n_components.
    """
    tsne = TSNE(n_components=n_components, perplexity=perplexity, random_state=42, init='random')
    vis_dims = tsne.fit_transform(arr)
    return vis_dims

def plot_embeddings(embed_arr_ls: List[np.ndarray], n_components: int, names: List[str], perplexity: int):
    """Plot transformed embedding vectors with predefined labels.

    Args:
        embed_arr_ls: a list of np.ndarray. Each np.ndarray is a matrix with embeddings corresponding to data examples.
        n_components: int. The number of components for tSNE.
        names: a list of str. The names of the data sources. The length of this list should be the same as the length of embed_arr_ls.

    Returns:
        None
    """
    vis_dims = transform_tSNE(np.concatenate(embed_arr_ls), n_components, perplexity)

    colors = ["red", "blue", "green", "orange", "purple"]
    colormap = matplotlib.colors.ListedColormap(colors)
    color_indices = []
    for label in range(len(embed_arr_ls)):
        color_indices += [label] * len(embed_arr_ls[label])
    assert (len(vis_dims) == len(color_indices))
    
    x = [x for x,y in vis_dims]
    y = [y for x,y in vis_dims]

    fig, ax = plt.subplots()
    ax.scatter(x, y, c=color_indices, cmap=colormap, alpha=0.3)
    
    for label in range(len(embed_arr_ls)):
        color = colors[label]
        label_indices = [i for i, value in enumerate(color_indices) if value == label]
        avg_x = np.array(x)[label_indices].mean()
        avg_y = np.array(y)[label_indices].mean()
        ax.scatter(avg_x, avg_y, marker='x', color=color, s=100, label=names[label])

    ax.legend()
    plt.title("RedPajama-v1 sample data visualized in language using t-SNE")
    plt.show()

# Define values.
data_dir =  # where the sample dataset is downloaded.
model = "togethercomputer/m2-bert-80M-8k-retrieval" # You can use a different model.
context_length = 8192
num_samples = 50
filenames = ["arxiv_sample.jsonl", "book_sample.jsonl", "c4_sample.jsonl", "github_sample.jsonl", "stackexchange_sample.jsonl"]

embedding_ls = []
for f_name in filenames:
    embed_arr = sample_and_extract_embeddings(
    	os.path.join(data_dir, f_name),
model_api_string=model,
num_samples=num_samples,
context_length=context_length
    )
embedding_ls.append(embed_arr)
plot_embeddings(embedding_ls, n_components=2, names=filenames, perplextiy=15)