File size: 2,810 Bytes
8e874f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """
OpenAI Embeddings model for QAFD-RAG
"""
from typing import List, Optional
import numpy as np
import logging
import os
logger = logging.getLogger(__name__)
class OpenAIEmbeddingModel:
"""OpenAI embedding model - standalone version"""
def __init__(self, global_config, embedding_model_name: Optional[str] = None, api_key: Optional[str] = None):
self.global_config = global_config
self.embedding_model_name = embedding_model_name or global_config.embedding_model_name
# ✅ FIX: Try explicit api_key parameter first, then environment variable
self.api_key = api_key or os.environ.get("OPENAI_API_KEY", "")
if not self.api_key:
raise ValueError(
"OPENAI_API_KEY not found in environment variables. "
"Please set it with: export OPENAI_API_KEY='your-key-here'"
)
# Determine dimensions based on model
if "text-embedding-3-large" in self.embedding_model_name:
self.embedding_dim = 3072 # Default for 3-large
elif "text-embedding-3-small" in self.embedding_model_name:
self.embedding_dim = 1536
elif "text-embedding-ada-002" in self.embedding_model_name:
self.embedding_dim = 1536
else:
self.embedding_dim = 1536 # Default
# Allow dimension override for 3-large (can reduce to 256-3072)
if hasattr(global_config, 'embedding_dimensions'):
self.embedding_dim = global_config.embedding_dimensions
self.batch_size = getattr(global_config, 'embedding_batch_size', 100)
self.max_seq_len = getattr(global_config, 'embedding_max_seq_len', 8191)
logger.info(f"Initializing OpenAI: {self.embedding_model_name}")
logger.info(f"✅ OpenAI model ready: {self.embedding_dim}-dim")
def batch_encode(self, texts: List[str], **kwargs) -> np.ndarray:
"""Encode texts to embeddings using OpenAI API"""
from openai import OpenAI
if isinstance(texts, str):
texts = [texts]
client = OpenAI(api_key=self.api_key)
# Build request params
request_params = {
"model": self.embedding_model_name,
"input": texts,
"encoding_format": "float",
}
# Add dimensions param only for text-embedding-3 models
if "text-embedding-3" in self.embedding_model_name:
request_params["dimensions"] = self.embedding_dim
# Encode with OpenAI
response = client.embeddings.create(**request_params)
# Extract embeddings
embeddings = np.array([item.embedding for item in response.data])
return embeddings |