Add experimental co.summarize API (#145)

cohere/__init__.py

@@ -10,6 +10,7 @@
 EMBED_URL = 'embed'
 FEEDBACK_URL = 'feedback'
 GENERATE_URL = 'generate'
+SUMMARIZE_URL = 'summarize'
 
 CHECK_API_KEY_URL = 'check-api-key'
 TOKENIZE_URL = 'tokenize'

cohere/client.py

@@ -21,6 +21,7 @@
 from cohere.feedback import Feedback
 from cohere.generation import Generations
 from cohere.tokenize import Tokens
+from cohere.summarize import SummarizeResponse
 
 use_xhr_client = False
 try:
@@ -205,6 +206,55 @@ def classify(self,
 
         return Classifications(classifications)
 
+    def summarize(self, text: str, model: str = None, length: str = None, format: str = None, temperature: float = None,
+                  additional_instruction: str = None) -> SummarizeResponse:
+        """Return a generated summary of the specified length for the provided text.
+
+        Args:
+            text (str): Text to summarize.
+            model (str): (Optional) ID of the model.
+            length (str): (Optional) One of {"short", "medium", "long"}, defaults to "medium". \
+                Controls the length of the summary.
+            format (str): (Optional) One of {"paragraph", "bullets"}, defaults to "bullets". \
+                Controls the format of the summary.
+            temperature (float): Ranges from 0 to 5. Controls the randomness of the output. \
+                Lower values tend to generate more “predictable” output, while higher values \
+                tend to generate more “creative” output. The sweet spot is typically between 0 and 1.
+            additional_instruction (str): (Optional) Modifier for the underlying prompt, must \
+                complete the sentence "Generate a summary _".
+
+        Example:
+        ```
+        res = co.summarize(text="Stock market report for today...")
+        print(res.summary)
+        ```
+
+        Example:
+        ```
+        res = co.summarize(
+            text="Stock market report for today...",
+            model="summarize-xlarge",
+            length="long",
+            format="bullets",
+            temperature=0.9,
+            additional_instruction="focusing on the highest performing stocks")
+        print(res.summary)
+        ```
+        """
+        json_body = {
+            'model': model,
+            'text': text,
+            'length': length,
+            'format': format,
+            'temperature': temperature,
+            'additional_instruction': additional_instruction,
+        }
+        # remove None values from the dict
+        json_body = {k: v for k, v in json_body.items() if v is not None}
+        response = self.__request(cohere.SUMMARIZE_URL, json=json_body)
+
+        return SummarizeResponse(id=response["id"], summary=response["summary"])
+
     def batch_tokenize(self, texts: List[str]) -> List[Tokens]:
         return [self.tokenize(t) for t in texts]
 

cohere/summarize.py

@@ -0,0 +1,24 @@
+from typing import NamedTuple
+
+SummarizeResponse = NamedTuple("SummarizeResponse", [("id", str), ("summary", float)])
+SummarizeResponse.__doc__ = """
+Returned by co.summarize, which generates a summary of the specified length for the provided text.
+
+Example:
+```
+res = co.summarize(text="Stock market report for today...")
+print(res.summary)
+```
+
+Example:
+```
+res = co.summarize(
+    text="Stock market report for today...",
+    model="summarize-xlarge",
+    length="long",
+    format="bullets",
+    temperature=0.3,
+    additional_instruction="focusing on the highest performing stocks")
+print(res.summary)
+```
+"""

setup.py

@@ -24,7 +24,7 @@ def has_ext_modules(foo) -> bool:
 
 
 setuptools.setup(name='cohere',
-                 version='3.2.3',
+                 version='3.2.5',
                  author='1vn',
                  author_email='ivan@cohere.ai',
                  description='A Python library for the Cohere API',

tests/test_summarize.py

@@ -0,0 +1,29 @@
+import unittest
+
+from utils import get_api_key
+
+import cohere
+from cohere.summarize import SummarizeResponse
+from cohere.error import CohereError
+
+API_KEY = get_api_key()
+co = cohere.Client(API_KEY)
+
+text = """Leah LaBelle Vladowski was born on September 8, 1986, in Toronto, \
+Ontario, and raised in Seattle, Washington. Her parents, Anastasia and Troshan Vladowski, \
+are Bulgarian singers and her uncle made rock music in Bulgaria. Anastasia recorded pop music \
+and was in a group with Troshan, who was a founding member of Bulgaria's first rock band, \
+Srebyrnite grivni.[3] After defecting from Bulgaria during a 1979 tour,[1][3] LaBelle's parents \
+emigrated to Canada and later the United States, becoming naturalized citizens in both countries. \
+LaBelle grew up listening to music, including jazz and the Beatles, but felt the most connected to R&B"""
+
+
+class TestFeedback(unittest.TestCase):
+
+    def test_success(self):
+        res = co.summarize(text)
+        self.assertIsInstance(res, SummarizeResponse)
+
+    def error(self):
+        res = co.summarize(text, length="potato")
+        self.assertIsInstance(res, CohereError)