Embed API로 배치 임베딩 작업(Embed Jobs) 처리하기

대량의 텍스트 데이터를 효율적으로 처리하는 방법을 데이터셋 생성과 임베딩 작업 실행을 중심으로 알아볼 거예요. 이 가이드에서는 embed jobs 엔드포인트로 대량의 텍스트를 비동기적으로 임베딩하는 방법을 보여드릴게요.

이 가이드는 Embed Jobs API를 사용해요.

자세한 내용은 API 레퍼런스를 참고해 주세요.

Embed Jobs API는 embed v3.0 모델에서만 호환돼요.

이 가이드에서는 embed jobs 엔드포인트로 대량의 텍스트를 비동기적으로 임베딩하는 방법을 보여드릴게요. 이 엔드포인트의 기능을 설명하기 위해 위키백과 페이지와 그에 연결된 메타데이터로 이루어진 간단한 데이터셋을 사용해요. 검색의 end-to-end 예시를 보고 싶다면 이 노트북을 확인해 보세요.

출처: 문서

Embed Jobs API 사용 방법

Embed Jobs API는 대규모 정보 코퍼스에서 검색의 힘을 활용하고 싶은 사용자를 위해 설계됐어요. API를 통해 수십만 개의 문서(또는 청크)를 인코딩하는 일은 지루하고 느려서, 보통 여러분의 시스템과 우리 서버 사이에 수백만 건의 http 요청이 오가게 돼요. Embed Jobs API는 검증(validation), 스테이징(staging), 일괄 처리 최적화(batching)를 사용자 대신 처리해 주기 때문에 많은 수(100K 이상)의 문서를 인코딩하는 데 훨씬 더 적합해요. 또 Embed Jobs API는 결과를 호스팅된 데이터셋에 저장하므로 임베딩 결과를 로컬에 저장할 필요가 없어요.

Embed Jobs API는 Embed API와 함께 동작해요. 운영(production) 환경에서는 Embed Jobs로 코퍼스에 대한 대규모 주기적 업데이트를 준비하고, Embed로 실시간 쿼리와 더 작은 실시간 업데이트를 처리하는 식이에요.

Embed Jobs용 데이터셋 구성하기

Embed Jobs용 데이터셋을 만들려면 데이터셋 type을 embed-input으로 설정해야 해요. 파일의 스키마는 text:string처럼 생겼어요.

Embed Jobs API와 Datasets API는 keep_fields, optional_fields 두 필드를 통해 메타데이터를 다뤄요. create dataset 단계에서 보존하고 싶은 메타데이터 필드에 해당하는 문자열 목록인 keep_fields나 optional_fields 중 하나를 지정할 수 있어요. keep_fields는 더 엄격해서, 항목에 해당 필드가 없으면 검증이 실패해요. 반면 optional_fields는 빈 필드를 건너뛰고 검증을 통과시킬 수 있어요.

샘플 데이터셋 입력 형식

JSONL

{"wiki_id": 69407798, "url": "https://en.wikipedia.org/wiki?curid=69407798", "views": 5674.4492597435465, "langs": 38, "title": "Deaths in 2022", "text": "The following notable deaths occurred in 2022. Names are reported under the date of death, in alphabetical order. A typical entry reports information in the following sequence:", "paragraph_id": 0, "id": 0}
{"wiki_id": 3524766, "url": "https://en.wikipedia.org/wiki?curid=3524766", "views": 5409.5609619796405, "title": "YouTube", "text": "YouTube is a global online video sharing and social media platform headquartered in San Bruno, California. It was launched on February 14, 2005, by Steve Chen, Chad Hurley, and Jawed Karim. It is owned by Google, and is the second most visited website, after Google Search. YouTube has more than 2.5 billion monthly users who collectively watch more than one billion hours of videos each day. , videos were being uploaded at a rate of more than 500 hours of content per minute.", "paragraph_id": 0, "id": 1}

위 예시에서 보듯이, langs는 첫 번째 항목에는 있지만 두 번째 항목에는 없으므로 이는 유효한 create_dataset 호출이에요. wiki_id, url, views, title 필드는 두 JSON 모두에 존재해요.

PYTHON

# Upload a dataset for embed jobs
ds = co.datasets.create(
    name="sample_file",
    # insert your file path here - you can upload it on the right - we accept .csv and jsonl files
    data=open("embed_jobs_sample_data.jsonl", "rb"),
    keep_fields=["wiki_id", "url", "views", "title"],
    optional_fields=["langs"],
    type="embed-input",
)

# wait for the dataset to finish validation

print(co.wait(ds))

cURL

curl --request POST \
  --url https://api.cohere.ai/v2/datasets \
  --header 'accept: application/json' \
  --header 'content-type: multipart/form-data' \
  --header "Authorization: bearer ***" \
  --form 'name=sample_file' \
  --form 'type=embed-input' \
  --form 'keep_fields=["wiki_id","url","views","title"]' \
  --form 'optional_fields=["langs"]' \
  --form 'data=@embed_jobs_sample_data.jsonl'

현재 데이터셋 엔드포인트는 .csv와 .jsonl 파일을 받아요. 두 경우 모두 text라는 필드나 text라는 헤더가 반드시 있어야 해요. 우리 저장소에서 유효한 jsonl 파일 예시와 유효한 csv 파일 예시를 확인할 수 있어요.

1. 데이터셋 업로드하기

Embed Jobs API는 dataset ID를 입력으로 받아요. dataset_type="embed-input"으로 로컬 파일을 Datasets API에 업로드하면 임베딩을 위한 데이터 검증이 이뤄져요. 데이터셋에는 text 필드가 있어야 해요. 현재 지원하는 입력 파일 형식은 .csv와 .jsonl이에요. 이 과정을 보여주는 코드 조각은 다음과 같아요.

PYTHON

import cohere

co = cohere.ClientV2(api_key="<YOUR API KEY>")

input_dataset = co.datasets.create(
    name="your_file_name",
    data=open("/content/your_file_path", "rb"),
    type="embed-input",
)

# block on server-side validation

print(co.wait(input_dataset))

cURL

curl --request POST \
  --url https://api.cohere.ai/v2/datasets \
  --header 'accept: application/json' \
  --header 'content-type: multipart/form-data' \
  --header "Authorization: bearer ***" \
  --form 'name=your_file_name' \
  --form 'type=embed-input' \
  --form 'data=@/content/your_file_path'

데이터셋을 업로드하면 다음과 같은 응답을 받게 돼요.

Text

uploading file, starting validation...

데이터셋이 업로드되고 검증되면 다음과 같은 응답을 받게 돼요.

TEXT

sample-file-m613zv was uploaded

데이터셋에서 검증 오류가 발생하면 datasets 페이지의 데이터셋 검증 오류 섹션을 참고해서 문제를 해결해 주세요.

2. Embed Job 시작하기

이제 데이터셋을 임베딩할 준비가 됐어요. 이를 보여주는 코드 조각은 다음과 같아요.

PYTHON

embed_job_response = co.embed_jobs.create(
    dataset_id=input_dataset.id,
    input_type="search_document",
    model="embed-english-v3.0",
    embedding_types=["float"],
    truncate="END",
)

# block until the job is complete

embed_job = co.wait(embed_job_response)

cURL

curl --request POST \
  --url https://api.cohere.ai/v2/embed-jobs \
  --header 'accept: application/json' \
  --header 'content-type: application/json' \
  --header "Authorization: bearer ***" \
  --data '{
    "dataset_id": "<YOUR_DATASET_ID>",
    "input_type": "search_document",
    "model": "embed-english-v3.0",
    "embedding_types": ["float"],
    "truncate": "END"
  }'

이 임베딩들 위에서 검색을 수행하고 싶고, 그것들이 우리의 지식 기반을 구성한다고 생각할 수 있으므로 input_type='search_document'로 설정했어요.

3. Embed Job 결과 저장 또는 조회하기

Embed Jobs의 출력은 데이터셋 객체인데, 이를 다운로드하거나 원하는 데이터베이스로 직접 전달(piping)할 수 있어요.

PYTHON

output_dataset_response = co.datasets.get(
    id=embed_job.output_dataset_id
)
output_dataset = output_dataset_response.dataset
co.utils.save_dataset(
    dataset=output_dataset,
    filepath="/content/embed_job_output.csv",
    format="csv",
)

cURL

curl --request GET \
  --url https://api.cohere.ai/v2/datasets/<DATASET_ID> \
  --header 'accept: application/json' \
  --header "Authorization: bearer ***" \

대신 데이터셋을 다운스트림 함수에 넘기고 싶다면 다음과 같이 할 수 있어요.

PYTHON

output_dataset_response = co.datasets.get(
    id=embed_job.output_dataset_id
)
output_dataset = output_dataset_response.dataset
results = []
for record in output_dataset:
    results.append(record)

샘플 출력

Embed Jobs API는 데이터셋의 원래 순서를 유지하고, 출력 데이터는 text: string, embedding: list of floats 스키마를 따르며, 임베딩 목록의 길이는 선택한 모델에 따라 달라져요. 예를 들어 embed-v4.0은 선택에 따라 256, 512, 1024, 1536(기본값) 중 하나가 되고, embed-english-light-v3.0은 384 dimensions가 돼요.

데이터셋을 jsonl로 다운로드한 경우 출력은 다음과 같은 모습일 거예요.

JSON

{
  "text": "The following notable deaths occurred in 2022. Names are reported under the date of death, in alphabetical order......",
  "embeddings": {
    "float":[0.006572723388671875, 0.0090484619140625, -0.02142333984375,....],
    "int8":null,
    "uint8":null,
    "binary":null,
    "ubinary":null
  }
}

데이터셋 업로드 시 optional_fields나 keep_fields로 보존하도록 지정한 메타데이터가 있다면, embed jobs의 출력은 다음과 같이 보일 거예요.

JSON

{
  "text": "The following notable deaths occurred in 2022. Names are reported under the date of death, in alphabetical order......",
  "embeddings": {
    "float":[0.006572723388671875, 0.0090484619140625, -0.02142333984375,....],
    "int8":null,
    "uint8":null,
    "binary":null,
    "ubinary":null
  },
  "field_one": "some_meta_data",
  "field_two": "some_meta_data",
}

다음 단계

Pinecone의 serverless 오퍼링을 활용한 검색의 end-to-end 노트북을 확인해 보세요.

더 알아보기 (Learn more)

Cohere 문서