34 lines
1.3 KiB
Python
34 lines
1.3 KiB
Python
import os
|
|
from typing import List, Dict, Any
|
|
import time
|
|
import numpy as np
|
|
from openai import OpenAI
|
|
from config import MODEL_NAME, OPENAI_API_KEY_ENV, BATCH_SIZE
|
|
|
|
def get_client() -> OpenAI:
|
|
api_key = os.getenv(OPENAI_API_KEY_ENV)
|
|
if not api_key:
|
|
raise RuntimeError(f"Missing OpenAI API key env var: {OPENAI_API_KEY_ENV}")
|
|
return OpenAI(api_key=api_key)
|
|
|
|
def embed_texts(texts: List[str], client: OpenAI = None, model: str = MODEL_NAME, batch_size: int = BATCH_SIZE) -> List[List[float]]:
|
|
"""Embed a list of texts in batches. Returns list of vectors (list[float])."""
|
|
if client is None:
|
|
client = get_client()
|
|
|
|
vectors: List[List[float]] = []
|
|
total = len(texts)
|
|
for i in range(0, total, batch_size):
|
|
batch = texts[i:i+batch_size]
|
|
# Retry loop (simple backoff); you can replace with tenacity if desired.
|
|
for attempt in range(5):
|
|
try:
|
|
resp = client.embeddings.create(model=model, input=batch)
|
|
for item in resp.data:
|
|
vectors.append(item.embedding)
|
|
break
|
|
except Exception as e:
|
|
wait = 2 ** attempt
|
|
print(f"[WARN] Embedding batch {i//batch_size+1}/{(total+batch_size-1)//batch_size} failed: {e}. Retrying in {wait}s...")
|
|
time.sleep(wait)
|
|
return vectors |