約6万件の日本語テキストをベクトル化をGPUとCPU処理で比較してみた

# CUDA対応のPyTorchイメージ
FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime

ENV PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    TZ=Asia/Tokyo

WORKDIR /app

RUN pip install --no-cache-dir sentence-transformers boto3

COPY vectorize-simple.py .

CMD ["python", "-u", "vectorize-simple.py"]

# CPU専用のPyTorchイメージ
FROM pytorch/pytorch:2.5.1-cpu

ENV PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1 \
    TZ=Asia/Tokyo

WORKDIR /app

RUN pip install --no-cache-dir sentence-transformers boto3

COPY vectorize-simple.py .

CMD ["python", "-u", "vectorize-simple.py"]

import json
import os
from datetime import datetime
from sentence_transformers import SentenceTransformer
import boto3

# 環境変数から設定を取得
bucket = os.environ['S3_BUCKET']
input_key = os.environ['INPUT_KEY']
output_key = os.environ['OUTPUT_KEY']

s3 = boto3.client('s3')

# S3から入力ファイルをダウンロード
print('Downloading input file...')
s3.download_file(bucket, input_key, '/tmp/articles_ja.json')

# モデルロード（GPUがあれば自動的に使用される）
model_name = os.environ.get('MODEL_NAME', 'intfloat/multilingual-e5-large')
print(f'Loading model: {model_name}...')
model = SentenceTransformer(model_name)

# 記事データ読み込み
print('Loading articles...')
with open('/tmp/articles_ja.json', 'r') as f:
    articles = json.load(f)

print(f'Processing {len(articles)} articles...')

# ベクトル化処理（1件ずつ処理、100件ごとに進捗表示）
with open('/tmp/vectors_ja.jsonl', 'w') as f:
    for i, article in enumerate(articles):
        text = f"{article['title']} {article['summary']}"
        vector = model.encode(text).tolist()
        result = {
            'article_id': article['id'],
            'vector': vector,
            'dimension': len(vector),
            'model': model_name,
            'timestamp': datetime.utcnow().isoformat()
        }
        f.write(json.dumps(result, ensure_ascii=False) + '\n')
        if (i + 1) % 100 == 0:
            print(f'Processed {i + 1}/{len(articles)} articles')

# S3にアップロード
print('Uploading results...')
s3.upload_file('/tmp/vectors_ja.jsonl', bucket, output_key)
print('Done!')

import json, gzip, os, time
from sentence_transformers import SentenceTransformer
import boto3

BUCKET = os.environ["S3_BUCKET"]
INPUT_KEY = os.environ["INPUT_KEY"]
OUTPUT_KEY = os.environ["OUTPUT_KEY"]
MODEL_NAME = os.environ.get("MODEL_NAME", "intfloat/multilingual-e5-large")
BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "256"))
PREFIX = os.environ.get("TEXT_PREFIX", "")

s3 = boto3.client("s3")

print(f"Downloading s3://{BUCKET}/{INPUT_KEY}")
s3.download_file(BUCKET, INPUT_KEY, "/tmp/input.jsonl.gz")

print(f"Loading model: {MODEL_NAME}")
model = SentenceTransformer(MODEL_NAME)

print("Reading input...")
records = []
with gzip.open("/tmp/input.jsonl.gz", "rt") as f:
    for line in f:
        d = json.loads(line)
        rid = d["recordId"]
        text = d["modelInput"]["singleEmbeddingParams"]["text"]["value"]
        records.append((rid, PREFIX + text))

print(f"Processing {len(records)} records with batch_size={BATCH_SIZE}")
start = time.time()

texts = [r[1] for r in records]
vectors = model.encode(texts, batch_size=BATCH_SIZE, show_progress_bar=True)

elapsed = time.time() - start
rate = len(records) / (elapsed / 60)
print(f"Encoding done: {elapsed:.1f}s ({rate:.0f} records/min)")

print("Writing output...")
with open("/tmp/output.jsonl", "w") as f:
    for i, (rid, _) in enumerate(records):
        out = {"recordId": rid, "vector": vectors[i].tolist(), "dimension": len(vectors[i]), "model": MODEL_NAME}
        f.write(json.dumps(out, ensure_ascii=False) + "\n")

print(f"Uploading to s3://{BUCKET}/{OUTPUT_KEY}")
s3.upload_file("/tmp/output.jsonl", BUCKET, OUTPUT_KEY)
print(f"Done! {len(records)} records, {elapsed:.1f}s, {rate:.0f} records/min")

docker run --gpus all --rm \
  -e S3_BUCKET=your-bucket-name \
  -e INPUT_KEY=articles_ja.json \
  -e OUTPUT_KEY=vectors_ja_gpu_${TIMESTAMP}.jsonl \
  vectorize-gpu

docker run --rm \
  -e S3_BUCKET=your-bucket-name \
  -e INPUT_KEY=articles_ja.json \
  -e OUTPUT_KEY=vectors_ja_cpu_${TIMESTAMP}.jsonl \
  vectorize-cpu

FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime

# 日本語形態素解析用パッケージを追加
RUN pip install --no-cache-dir sentence-transformers boto3 protobuf fugashi unidic-lite

COPY vectorize-ruri-base.py /app/vectorize-ruri-base.py

WORKDIR /app

CMD ["python", "vectorize-ruri-base.py"]

FROM pytorch/pytorch:2.5.1-cpu

# 日本語形態素解析用パッケージを追加
RUN pip install --no-cache-dir sentence-transformers boto3 protobuf fugashi unidic-lite

COPY vectorize-ruri-base.py /app/vectorize-ruri-base.py

WORKDIR /app

CMD ["python", "vectorize-ruri-base.py"]

import json
import os
from datetime import datetime
from sentence_transformers import SentenceTransformer
import boto3

bucket = os.environ['S3_BUCKET']
input_key = os.environ['INPUT_KEY']
output_key = os.environ['OUTPUT_KEY']

s3 = boto3.client('s3')

print('Downloading input file...')
s3.download_file(bucket, input_key, '/tmp/articles_ja.json')

print('Loading model: cl-nagoya/ruri-base...')
model = SentenceTransformer('cl-nagoya/ruri-base')

print('Loading articles...')
with open('/tmp/articles_ja.json', 'r') as f:
    articles = json.load(f)

print(f'Processing {len(articles)} articles...')

with open('/tmp/vectors_ja.jsonl', 'w') as f:
    for i, article in enumerate(articles):
        # ruriモデルは "文章: " プレフィックスを推奨
        text = f"文章: {article['title']} {article['summary']}"
        vector = model.encode(text).tolist()
        result = {
            'article_id': article['id'],
            'vector': vector,
            'dimension': len(vector),
            'model': 'cl-nagoya/ruri-base',
            'timestamp': datetime.utcnow().isoformat()
        }
        f.write(json.dumps(result, ensure_ascii=False) + '\n')
        if (i + 1) % 100 == 0:
            print(f'Processed {i + 1}/{len(articles)} articles')

print('Uploading results...')
s3.upload_file('/tmp/vectors_ja.jsonl', bucket, output_key)
print('Done!')

docker run --gpus all --rm \
  -e S3_BUCKET=your-bucket-name \
  -e INPUT_KEY=articles_ja.json \
  -e OUTPUT_KEY=vectors_ja_ruri_gpu_${TIMESTAMP}.jsonl \
  vectorize-ruri-gpu

docker run --rm \
  -e S3_BUCKET=your-bucket-name \
  -e INPUT_KEY=articles_ja.json \
  -e OUTPUT_KEY=vectors_ja_ruri_cpu_${TIMESTAMP}.jsonl \
  vectorize-ruri-cpu