Text Classification

Among the most popular use cases for language embeddings is ‘text classification,’ in which different pieces of text — blog posts, lyrics, poems, headlines, etc. — are grouped based on their similarity, their sentiment, or some other property.

Here, we’ll discuss how to perform simple text classification tasks with Cohere’s classify endpoint, and provide links to more information on how to fine-tune this endpoint for more specialized work.

Few-Shot Classification with Cohere’s classify Endpoint

Generally, training a text classifier requires a tremendous amount of data. But with large language models, it’s now possible to create so-called ‘few shot’ classification models able to perform well after seeing a far smaller number of samples.

In the next few sections, we’ll create a sentiment analysis classifier to sort text into “positive,” “negative,” and “neutral” categories.

Setting up the SDK

First, let’s import the required tools and set up a Cohere client.

1
import cohere
2
from cohere import ClassifyExample

1
co = cohere.ClientV2("COHERE_API_KEY") # Your Cohere API key

Preparing the Data and Inputs

With the classify endpoint, you can create a text classifier with as few as two examples per class, and each example must contain the text itself and the corresponding label (i.e. class). So, if you have two classes you need a minimum of four examples, if you have three classes you need a minimum of six examples, and so on.

Here are examples, created as ClassifyExample objects:

1
examples = [ClassifyExample(text="I’m so proud of you", label="positive"), 
2
            ClassifyExample(text="What a great time to be alive", label="positive"), 
3
            ClassifyExample(text="That’s awesome work", label="positive"), 
4
            ClassifyExample(text="The service was amazing", label="positive"), 
5
            ClassifyExample(text="I love my family", label="positive"), 
6
            ClassifyExample(text="They don't care about me", label="negative"), 
7
            ClassifyExample(text="I hate this place", label="negative"), 
8
            ClassifyExample(text="The most ridiculous thing I've ever heard", label="negative"), 
9
            ClassifyExample(text="I am really frustrated", label="negative"), 
10
            ClassifyExample(text="This is so unfair", label="negative"),
11
            ClassifyExample(text="This made me think", label="neutral"), 
12
            ClassifyExample(text="The good old days", label="neutral"), 
13
            ClassifyExample(text="What's the difference", label="neutral"), 
14
            ClassifyExample(text="You can't ignore this", label="neutral"), 
15
            ClassifyExample(text="That's how I see it", label="neutral")]

Besides the examples, you’ll also need the ‘inputs,’ which are the strings of text you want the classifier to sort. Here are the ones we’ll be using:

1
inputs = ["Hello, world! What a beautiful day",
2
          "It was a great time with great people",
3
          "Great place to work",
4
          "That was a wonderful evening",
5
          "Maybe this is why",
6
          "Let's start again",
7
          "That's how I see it",
8
          "These are all facts",
9
          "This is the worst thing",
10
          "I cannot stand this any longer",
11
          "This is really annoying",
12
          "I am just plain fed up"]

Generate Predictions

Setting up the model is quite straightforward with the classify endpoint. We’ll use Cohere’s embed-english-v3.0 model, here’s what that looks like:

1
def classify_text(inputs, examples):
2
  
3
    """
4
    Classifies a list of input texts given the examples
5
    Arguments:
6
        model (str): identifier of the model
7
        inputs (list[str]): a list of input texts to be classified
8
        examples (list[Example]): a list of example texts and class labels
9
    Returns:
10
        classifications (list): each result contains the text, labels, and conf values
11
    """
12
    
13
    # Classify text by calling the Classify endpoint
14
    response = co.classify(
15
        model='embed-english-v3.0',
16
        inputs=inputs,
17
        examples=examples)
18
19
    classifications = response.classifications
20
21
    return classifications
22
23
# Classify the inputs
24
predictions = classify_text(inputs, examples)
25
26
print(predictions) 

Here’s a sample output returned (note that this output has been truncated to make it easier to read, you’ll get much more in return if you run the code yourself):

[ClassifyResponseClassificationsItem(id='9df6628d-57b2-414c-837e-c8a22f00d3db',
                                     input='hello, world! what a beautiful day',
                                     prediction='positive',
                                     predictions=['positive'],
                                     confidence=0.40137812,
                                     confidences=[0.40137812],
                                     labels={'negative': ClassifyResponseClassificationsItemLabelsValue(confidence=0.23582731),
                                             'neutral': ClassifyResponseClassificationsItemLabelsValue(confidence=0.36279458),
                                             'positive': ClassifyResponseClassificationsItemLabelsValue(confidence=0.40137812)},
                                     classification_type='single-label'),
 ClassifyResponseClassificationsItem(id='ce2c3b0b-ce98-4905-9ef5-fc83c6848fc5',
                                     input='it was a great time with great people',
                                     prediction='positive',
                                     predictions=['positive'],
                                     confidence=0.49054274,
                                     confidences=[0.49054274],
                                     labels={'negative': ClassifyResponseClassificationsItemLabelsValue(confidence=0.19989403),
                                             'neutral': ClassifyResponseClassificationsItemLabelsValue(confidence=0.30956325),
                                             'positive': ClassifyResponseClassificationsItemLabelsValue(confidence=0.49054274)},
                                     classification_type='single-label')
 ....]

Most of this is pretty easy to understand, but there are a few things worth drawing attention to.

Besides returning the predicted class in the prediction field, the endpoint also returns the confidence value of the prediction, which varies between 0 (unconfident) and 1 (completely confident).

Also, these confidence values are split among the classes; since we’re using three, the confidence values for the “positive,” “negative,” and “neutral” classes must add up to a total of 1.

Under the hood, the classifier selects the class with the highest confidence value as the “predicted class.” A high confidence value for the predicted class therefore indicates that the model is very confident of its prediction, and vice versa.

What If I Need to Fine-Tune the classify endpoint?

Cohere has dedicated documentation on fine-tuning the classify endpoint for bespoke tasks. You can also read this blog post, which works out a detailed example.