Long-Form Text Strategies with Cohere

1
%%capture
2
!pip install cohere
3
!pip install python-dotenv
4
!pip install tokenizers
5
!pip install langchain
6
!pip install nltk
7
!pip install networkx
8
!pip install pypdf2

1
import os
2
import requests
3
from collections import deque
4
from typing import List, Tuple
5
6
import cohere
7
8
import numpy as np
9
10
import PyPDF2
11
from dotenv import load_dotenv
12
13
from tokenizers import Tokenizer
14
15
import nltk
16
nltk.download('punkt')  # Download the necessary data for sentence tokenization
17
from nltk.tokenize import sent_tokenize
18
19
import networkx as nx
20
from getpass import getpass

[nltk_data] Downloading package punkt to
[nltk_data]     /home/anna_cohere_com/nltk_data...
[nltk_data]   Package punkt is already up-to-date!

1
# Set up Cohere client
2
co_model = 'command-a-03-2025'
3
co_api_key = getpass("Enter your Cohere API key: ")
4
co = cohere.Client(api_key=co_api_key)

1
def load_long_pdf(file_path):
2
    """
3
    Load a long PDF file and extract its text content.
4
5
    Args:
6
        file_path (str): The path to the PDF file.
7
8
    Returns:
9
        str: The extracted text content of the PDF file.
10
    """
11
    with open(file_path, 'rb') as file:
12
        pdf_reader = PyPDF2.PdfReader(file)
13
        num_pages = len(pdf_reader.pages)
14
        full_text = ''
15
        for page_num in range(num_pages):
16
            page = pdf_reader.pages[page_num]
17
            full_text += page.extract_text()
18
    return full_text
19
20
def save_pdf_from_url(pdf_url, save_path):
21
    try:
22
        # Send a GET request to the PDF URL
23
        response = requests.get(pdf_url, stream=True)
24
        response.raise_for_status()  # Raise an exception for HTTP errors
25
26
        # Open the local file for writing in binary mode
27
        with open(save_path, 'wb') as file:
28
            # Write the content of the response to the local file
29
            for chunk in response.iter_content(chunk_size=8192):
30
                file.write(chunk)
31
32
        print(f"PDF saved successfully to '{save_path}'")
33
    except requests.exceptions.RequestException as e:
34
        print(f"Error downloading PDF: {e}")

1
# Download the PDF file from the URL
2
pdf_url = 'https://data.consilium.europa.eu/doc/document/ST-5662-2024-INIT/en/pdf'
3
save_path = 'example.pdf'
4
save_pdf_from_url(pdf_url, save_path)
5
6
# Load the PDF file and extract its text content
7
long_text = load_long_pdf(save_path)
8
long_text = long_text.replace('\n', ' ')
9
10
# Print the length of the document
11
print("Document length - #tokens:", len(co.tokenize(text=long_text, model=co_model).tokens))

PDF saved successfully to 'example.pdf'
Document length - #tokens: 134184

1
def generate_response(message, max_tokens=300, temperature=0.2, k=0):
2
  """
3
  A wrapper around the Cohere API to generate a response based on a given prompt.
4
5
  Args:
6
    messsage (str): The input message for generating the response.
7
    max_tokens (int, optional): The maximum number of tokens in the generated response. Defaults to 300.
8
    temperature (float, optional): Controls the randomness of the generated response. Higher values (e.g., 1.0) make the output more random, while lower values (e.g., 0.2) make it more deterministic. Defaults to 0.2.
9
    k (int, optional): Controls the diversity of the generated response. Higher values (e.g., 5) make the output more diverse, while lower values (e.g., 0) make it more focused. Defaults to 0.
10
11
  Returns:
12
    str: The generated response.
13
14
  """
15
  response = co.chat(
16
    model = co_model,
17
    message=message,
18
    max_tokens=max_tokens,
19
    temperature=temperature
20
    )
21
  return response.text

1
# Example summary prompt.
2
prompt_template = """
3
## Instruction
4
Summarize the following Document in 3-5 sentences. Only answer based on the information provided in the document.
5
6
## Document
7
{document}
8
9
## Summary
10
""".strip()

1
prompt = prompt_template.format(document=long_text)
2
# print(generate_response(message=prompt))

1
# The new Cohere model has a context limit of 128k tokens. However, for the purpose of this exercise, we will assume a smaller context window.
2
# Employing a smaller context window also has the additional benefit of reducing the cost per request, especially if billed by the number of tokens.
3
4
MAX_TOKENS = 40000
5
6
def truncate(long: str, max_tokens: int) -> str:
7
    """
8
    Shortens `long` by brutally truncating it to the first `max_tokens` tokens.
9
    This can break up sentences, passages, etc.
10
    """
11
12
    tokenized = co.tokenize(text=long, model=co_model).token_strings
13
    truncated = tokenized[:max_tokens]
14
    short = "".join(truncated)
15
    return short

1
short_text = truncate(long_text, MAX_TOKENS)
2
3
prompt = prompt_template.format(document=short_text)
4
print(generate_response(message=prompt))

1
def split_text_into_sentences(text) -> List[str]:
2
    """
3
    Split the input text into a list of sentences.
4
    """
5
    sentences = sent_tokenize(text)
6
7
    return sentences
8
9
def group_sentences_into_passages(sentence_list, n_sentences_per_passage=5):
10
    """
11
    Group sentences into passages of n_sentences sentences.
12
    """
13
    passages = []
14
    passage = ""
15
    for i, sentence in enumerate(sentence_list):
16
        passage += sentence + " "
17
        if (i + 1) % n_sentences_per_passage == 0:
18
            passages.append(passage)
19
            passage = ""
20
    return passages
21
22
def build_simple_chunks(text, n_sentences=5):
23
    """
24
    Build chunks of text from the input text.
25
    """
26
    sentences = split_text_into_sentences(text)
27
    chunks = group_sentences_into_passages(sentences, n_sentences_per_passage=n_sentences)
28
    return chunks

1
sentences = split_text_into_sentences(long_text)
2
passages = group_sentences_into_passages(sentences, n_sentences_per_passage=5)
3
print('Example sentence:', np.random.choice(np.asarray(sentences), size=1, replace=False))
4
print()
5
print('Example passage:', np.random.choice(np.asarray(passages), size=1, replace=False))

Example sentence: ['The European Data Protection Supervisor may also establish an AI regulatory sandbox for  the EU institutions, bodies and agencies and exercise the roles and the tasks of national  competent authorities in accordance with this chapter.']
Example passage: ['This flexibility could mean, for example a decision  by the provider to integrate a part of the necessary testing and reporting processes,  information and documentation required under this Regulation into already existing  documentation and procedu res required under the existing Union harmonisation legislation  listed in Annex II, Section A. This however should not in any way undermine the  obligation of the provider to comply with all the applicable requirements. (42a)   The risk management system shou ld consist of a continuous, iterative process that is  planned and run throughout the entire lifecycle of a high - risk AI system. This process  should be aimed at identifying and mitigating the relevant risks of artificial intelligence  systems on health, safe ty and fundamental rights. The risk management system should be  regularly reviewed and updated to ensure its continuing effectiveness, as well as  justification and documentation of any significant decisions and actions taken subject to  this Regulation. ']

1
def _add_chunks_by_priority(
2
    chunks: List[str],
3
    idcs_sorted_by_priority: List[int],
4
    max_tokens: int,
5
) -> List[Tuple[int, str]]:
6
    """
7
    Given chunks of text and their indices sorted by priority (highest priority first), this function
8
    fills the model context window with as many highest-priority chunks as possible.
9
10
    The output is a list of (index, chunk) pairs, ordered by priority. To stitch back the chunks into
11
    a cohesive text that preserves chronological order, sort the output on its index.
12
    """
13
14
    selected = []
15
    num_tokens = 0
16
    idcs_queue = deque(idcs_sorted_by_priority)
17
18
    while num_tokens < max_tokens and len(idcs_queue) > 0:
19
        next_idx = idcs_queue.popleft()
20
        num_tokens += len(co.tokenize(text=chunks[next_idx], model=co_model).tokens)
21
        # keep index and chunk, to reorder chronologically
22
        selected.append((next_idx, chunks[next_idx]))
23
    if num_tokens > max_tokens:
24
        selected.pop()
25
26
    return selected
27
28
def query_based_retrieval(
29
    long: str,
30
    max_tokens: int,
31
    query: str,
32
    n_setences_per_passage: int = 5,
33
) -> str:
34
    """
35
    Performs query-based retrieval on a long text document.
36
    """
37
    # 1. Chunk text into units
38
    chunks = build_simple_chunks(long, n_setences_per_passage)
39
40
    # 2. Use co.rerank to rank chunks vs. query
41
    chunks_reranked = co.rerank(query=query, documents=chunks, model="rerank-english-v3.0")
42
    idcs_sorted_by_relevance = [
43
        chunk.index for chunk in sorted(chunks_reranked.results, key=lambda c: c.relevance_score, reverse=True)
44
    ]
45
46
    # 3. Add chunks back in order of relevance
47
    selected = _add_chunks_by_priority(chunks, idcs_sorted_by_relevance, max_tokens)
48
49
    # 4. Put condensed text back in original order
50
    separator = " "
51
    short = separator.join([chunk for index, chunk in sorted(selected, key=lambda item: item[0], reverse=False)])
52
    return short

1
# Example prompt
2
prompt_template = """
3
## Instruction
4
{query}
5
6
## Document
7
{document}
8
9
## Answer
10
""".strip()

1
query = "What does the report say about biometric identification? Answer only based on the document."
2
short_text = query_based_retrieval(long_text, MAX_TOKENS, query)
3
prompt = prompt_template.format(query=query, document=short_text)
4
print(generate_response(message=prompt, max_tokens=300))

The report outlines several key points regarding biometric identification within the context of the proposed Artificial Intelligence Act:
1. **Prohibition of Real-Time Biometric Identification in Public Spaces**: The report proposes a ban on real-time biometric identification by law enforcement authorities in publicly accessible spaces, with specific exceptions. These exceptions are detailed in Article 5(1)(d) and are subject to safeguards, including monitoring, oversight, and limited reporting obligations at the EU level.
2. **Exceptions to the Prohibition**: The exceptions to the ban on real-time biometric identification include:
   - Search for victims of specific crimes (e.g., abduction, trafficking, sexual exploitation).
   - Prevention of imminent threats to life or physical safety, including terrorist attacks.
   - Localization or identification of suspects for serious criminal offenses (as defined in Annex IIa) punishable by a custodial sentence of at least four years.
3. **Safeguards and Conditions**: The use of real-time biometric identification systems in these exceptional cases must comply with specific safeguards and conditions, including:
   - A fundamental rights impact assessment.
   - Registration of the system in a database.
   - Prior authorization by a judicial or independent administrative authority, except in urgent situations where authorization can be sought within 24 hours.
   - Limitation to what is strictly necessary in terms of time, geography, and personal scope.
4. **Post-Remote Biometric Identification**: The use of post-remote biometric identification

1
def text_rank(text: str, max_tokens: int, n_setences_per_passage: int) -> str:
2
    """
3
    Shortens text by extracting key units of text from it based on their centrality.
4
    The output is the concatenation of those key units, in their original order.
5
    """
6
7
    # 1. Chunk text into units
8
    chunks = build_simple_chunks(text, n_setences_per_passage)
9
10
    # 2. Embed and construct similarity matrix
11
    embeddings = np.array(
12
        co.embed(
13
            texts=chunks,
14
            model="embed-v4.0",
15
            input_type="clustering",
16
        ).embeddings
17
    )
18
    similarities = np.dot(embeddings, embeddings.T)
19
20
    # 3. Compute centrality and sort sentences by centrality
21
    # Easiest to use networkx's `degree` function with similarity as weight
22
    g = nx.from_numpy_array(similarities, edge_attr="weight")
23
    centralities = g.degree(weight="weight")
24
    idcs_sorted_by_centrality = [node for node, degree in sorted(centralities, key=lambda item: item[1], reverse=True)]
25
26
    # 4. Add chunks back in order of centrality
27
    selected = _add_chunks_by_priority(chunks, idcs_sorted_by_centrality, max_tokens)
28
29
    # 5. Put condensed text back in original order
30
    short = " ".join([chunk for index, chunk in sorted(selected, key=lambda item: item[0], reverse=False)])
31
32
    return short

1
# Example summary prompt.
2
prompt_template = """
3
## Instruction
4
Summarize the following Document in 3-5 sentences. Only answer based on the information provided in the document.
5
6
## Document
7
{document}
8
9
## Summary
10
""".strip()

1
short_text = text_rank(long_text, MAX_TOKENS, 5)
2
prompt = prompt_template.format(document=short_text)
3
print(generate_response(message=prompt, max_tokens=600))

The document outlines the European Union's regulatory framework for artificial intelligence (AI) systems, focusing on high-risk AI applications. It establishes rules for placing AI systems on the market, including prohibitions on certain practices, requirements for high-risk systems, and transparency obligations. The regulation defines high-risk AI systems based on their intended use and potential risks to health, safety, and fundamental rights. Providers of high-risk AI systems must comply with specific requirements, such as risk management, data governance, and human oversight. The regulation also mandates conformity assessments, registration in an EU database, and post-market monitoring. It emphasizes the importance of AI literacy, prohibits manipulative or exploitative AI practices, and ensures compliance through market surveillance and enforcement mechanisms. Additionally, the regulation addresses general-purpose AI models, requiring providers to meet specific obligations, especially for models with systemic risks. The framework aims to promote trustworthy AI while safeguarding public interests and supporting innovation.