Merge branch 'main' into feature/extract-endpoint

.github/workflows/lint.yaml

@@ -1,9 +1,6 @@
 name: Lint
 
-on:
-  push:
-    paths:
-      - '**.py'
+on: push
 
 jobs:
   flake8:
@@ -21,6 +18,6 @@ jobs:
       - name: Run flake8
         uses: suo/flake8-github-action@releases/v1
         with:
-          checkName: 'flake8'
+          checkName: "flake8"
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/test.yaml

@@ -10,4 +10,4 @@ jobs:
           python-version: "3.6"
       - env:
           CO_API_KEY: ${{ secrets.CO_API_PROD_KEY }}
-        run: pip install requests && python -m unittest tests/tests.py
+        run: pip install requests && python -m unittest discover tests

tests/test_classify.py

@@ -0,0 +1,108 @@
+import unittest
+import cohere
+from cohere.classify import Example
+from utils import get_api_key
+
+co = cohere.Client(get_api_key())
+
+
+class TestClassify(unittest.TestCase):
+    def test_success(self):
+        prediction = co.classify('medium', ['purple'], [
+            Example('apple', 'fruit'),
+            Example('banana', 'fruit'),
+            Example('cherry', 'fruit'),
+            Example('watermelon', 'fruit'),
+            Example('kiwi', 'fruit'),
+            Example('red', 'color'),
+            Example('blue', 'color'),
+            Example('green', 'color'),
+            Example('yellow', 'color'),
+            Example('magenta', 'color')
+        ])
+        self.assertIsInstance(prediction.classifications, list)
+        self.assertIsInstance(prediction.classifications[0].input, str)
+        self.assertIsInstance(prediction.classifications[0].prediction, str)
+        self.assertIsInstance(prediction.classifications[0].confidence[0].confidence, (int, float))
+        self.assertIsInstance(prediction.classifications[0].confidence[0].label, str)
+        self.assertIsInstance(prediction.classifications[0].confidence[1].confidence, (int, float))
+        self.assertIsInstance(prediction.classifications[0].confidence[1].label, str)
+        self.assertEqual(len(prediction.classifications), 1)
+        self.assertEqual(prediction.classifications[0].prediction, 'color')
+
+    def test_empty_inputs(self):
+        with self.assertRaises(cohere.CohereError):
+            co.classify(
+                'medium', [], [
+                    Example('apple', 'fruit'),
+                    Example('banana', 'fruit'),
+                    Example('cherry', 'fruit'),
+                    Example('watermelon', 'fruit'),
+                    Example('kiwi', 'fruit'),
+
+                    Example('red', 'color'),
+                    Example('blue', 'color'),
+                    Example('green', 'color'),
+                    Example('yellow', 'color'),
+                    Example('magenta', 'color')])
+
+    def test_success_multi_input(self):
+        prediction = co.classify('medium', ['purple', 'mango'], [
+            Example('apple', 'fruit'),
+            Example('banana', 'fruit'),
+            Example('cherry', 'fruit'),
+            Example('watermelon', 'fruit'),
+            Example('kiwi', 'fruit'),
+
+            Example('red', 'color'),
+            Example('blue', 'color'),
+            Example('green', 'color'),
+            Example('yellow', 'color'),
+            Example('magenta', 'color')])
+        self.assertEqual(prediction.classifications[0].prediction, 'color')
+        self.assertEqual(prediction.classifications[1].prediction, 'fruit')
+        self.assertEqual(len(prediction.classifications), 2)
+
+    def test_success_all_fields(self):
+        prediction = co.classify('medium', ['mango', 'purple'], [
+            Example('apple', 'fruit'),
+            Example('banana', 'fruit'),
+            Example('cherry', 'fruit'),
+            Example('watermelon', 'fruit'),
+            Example('kiwi', 'fruit'),
+
+            Example('red', 'color'),
+            Example('blue', 'color'),
+            Example('green', 'color'),
+            Example('yellow', 'color'),
+            Example('magenta', 'color')
+        ], 'this is a classifier to determine if a word is a fruit of a color', 'This is a')
+        self.assertEqual(prediction.classifications[0].prediction, 'fruit')
+        self.assertEqual(prediction.classifications[1].prediction, 'color')
+
+    def test_string_repr(self):
+        prediction = co.classify('medium', ['purple'], [
+            Example('apple', 'fruit'),
+            Example('banana', 'fruit'),
+            Example('cherry', 'fruit'),
+            Example('watermelon', 'fruit'),
+            Example('kiwi', 'fruit'),
+            Example('red', 'color'),
+            Example('blue', 'color'),
+            Example('green', 'color'),
+            Example('yellow', 'color'),
+            Example('magenta', 'color')
+        ])
+        self.assertEqual(repr(prediction), '''cohere.Classifications {
+\tclassifications: [cohere.Classification {
+\tinput: purple
+\tprediction: color
+\tconfidence: [cohere.Confidence {
+\tlabel: fruit
+\tconfidence: 0
+}, cohere.Confidence {
+\tlabel: color
+\tconfidence: 1
+}]
+}]
+}''')

tests/test_embed.py

@@ -0,0 +1,93 @@
+import unittest
+import random
+import string
+import cohere
+from utils import get_api_key
+
+co = cohere.Client(get_api_key())
+
+
+def random_word():
+    return ''.join(random.choice(string.ascii_lowercase) for _ in range(10))
+
+
+def random_sentence(num_words):
+    sentence = ''
+
+    for _ in range(num_words):
+        sentence += random_word() + ' '
+
+    return sentence
+
+
+def random_texts(num_texts, num_words_per_sentence=50):
+    arr = []
+
+    for _ in range(num_texts):
+        arr.append(random_sentence(num_words_per_sentence))
+
+    return arr
+
+
+class TestEmbed(unittest.TestCase):
+    def test_success(self):
+        prediction = co.embed(
+            model='small',
+            texts=['co:here', 'cohere'])
+        self.assertEqual(len(prediction.embeddings), 2)
+        self.assertIsInstance(prediction.embeddings[0], list)
+        self.assertIsInstance(prediction.embeddings[1], list)
+        self.assertEqual(len(prediction.embeddings[0]), 1024)
+        self.assertEqual(len(prediction.embeddings[1]), 1024)
+
+    def test_success_multiple_batches(self):
+        prediction = co.embed(
+            model='small',
+            texts=['co:here', 'cohere', 'embed', 'python', 'golang', 'typescript', 'rust?', 'ai', 'nlp', 'neural'])
+        self.assertEqual(len(prediction.embeddings), 10)
+        for embed in prediction.embeddings:
+            self.assertIsInstance(embed, list)
+            self.assertEqual(len(embed), 1024)
+
+    def test_success_longer_multiple_batches_unaligned_batch(self):
+        prediction = co.embed(
+            model='small',
+            texts=[
+                'co:here', 'cohere', 'embed', 'python', 'golang',
+                'typescript', 'rust?',  'ai',  'nlp', 'neural', 'nets'
+            ])
+        self.assertEqual(len(prediction.embeddings), 11)
+        for embed in prediction.embeddings:
+            self.assertIsInstance(embed, list)
+            self.assertEqual(len(embed), 1024)
+
+    def test_success_longer_multiple_batches(self):
+        prediction = co.embed(
+            model='small',
+            texts=['co:here', 'cohere', 'embed', 'python', 'golang'] * 200)
+        self.assertEqual(len(prediction.embeddings), 200*5)
+        for embed in prediction.embeddings:
+            self.assertIsInstance(embed, list)
+            self.assertEqual(len(embed), 1024)
+
+    def test_success_multiple_batches_in_order(self):
+        textAll = []
+        predictionsExpected = []
+
+        for _ in range(3):
+            text_batch = random_texts(cohere.COHERE_EMBED_BATCH_SIZE)
+            prediction = co.embed(
+                model='small',
+                texts=text_batch)
+            textAll.extend(text_batch)
+            predictionsExpected.extend(prediction)
+        predictionsActual = co.embed(model='small', texts=textAll)
+        for predictionExpected, predictionActual in zip(predictionsExpected, list(predictionsActual)):
+            for elementExpected, elementAcutal in zip(predictionExpected, predictionActual):
+                self.assertAlmostEqual(elementExpected, elementAcutal, places=1)
+
+    def test_invalid_texts(self):
+        with self.assertRaises(cohere.CohereError):
+            co.embed(
+                model='small',
+                texts=[''])

tests/test_generate.py

@@ -0,0 +1,70 @@
+import unittest
+import cohere
+from utils import get_api_key
+
+API_KEY = get_api_key()
+co = cohere.Client(API_KEY)
+
+
+class TestGenerate(unittest.TestCase):
+    def test_success(self):
+        prediction = co.generate(
+            model='small',
+            prompt='co:here',
+            max_tokens=1)
+        self.assertIsInstance(prediction.generations[0].text, str)
+        self.assertIsNone(prediction.generations[0].token_likelihoods)
+        self.assertEqual(prediction.return_likelihoods, 'NONE')
+
+    def test_return_likelihoods_generation(self):
+        prediction = co.generate(
+            model='small',
+            prompt='co:here',
+            max_tokens=1,
+            return_likelihoods='GENERATION')
+        self.assertTrue(prediction.generations[0].token_likelihoods)
+        self.assertTrue(prediction.generations[0].token_likelihoods[0].token)
+        self.assertIsNotNone(prediction.generations[0].likelihood)
+        self.assertEqual(prediction.return_likelihoods, 'GENERATION')
+
+    def test_return_likelihoods_all(self):
+        prediction = co.generate(
+            model='small',
+            prompt='hi',
+            max_tokens=1,
+            return_likelihoods='ALL')
+        self.assertEqual(len(prediction.generations[0].token_likelihoods), 2)
+        self.assertIsNotNone(prediction.generations[0].likelihood)
+        self.assertEqual(prediction.return_likelihoods, 'ALL')
+
+    def test_invalid_temp(self):
+        with self.assertRaises(cohere.CohereError):
+            co.generate(
+                model='large',
+                prompt='hi',
+                max_tokens=1,
+                temperature=-1)
+
+    def test_invalid_model(self):
+        with self.assertRaises(cohere.CohereError):
+            co.generate(
+                model='this-better-not-exist',
+                prompt='co:here',
+                max_tokens=1)
+
+    def test_no_version_works(self):
+        cohere.Client(API_KEY).generate(
+            model='small',
+            prompt='co:here',
+            max_tokens=1)
+
+    def test_invalid_version_fails(self):
+        with self.assertRaises(cohere.CohereError):
+            cohere.Client(API_KEY, 'fake').generate(
+                model='small',
+                prompt='co:here',
+                max_tokens=1)
+
+    def test_invalid_key(self):
+        with self.assertRaises(cohere.CohereError):
+            _ = cohere.Client('invalid')

tests/test_tokenize.py

@@ -0,0 +1,17 @@
+import unittest
+import cohere
+from utils import get_api_key
+
+co = cohere.Client(get_api_key())
+
+
+class TestTokenize(unittest.TestCase):
+    def test_success(self):
+        tokens = co.tokenize('medium', 'tokenize me!')
+        self.assertIsInstance(tokens.tokens, list)
+        self.assertIsInstance(tokens.length, int)
+        self.assertEqual(tokens.length, len(tokens))
+
+    def test_invalid_text(self):
+        with self.assertRaises(cohere.CohereError):
+            co.tokenize(model='medium', text='')