Spaces:
Sleeping
Sleeping
File size: 1,642 Bytes
26fe9a7 473fd36 26fe9a7 |
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 |
import os
from typing import List
import numpy as np
from sentence_transformers import SentenceTransformer
from openai import OpenAI
class Embedder:
def __init__(self, model_name: str = "all-MiniLM-L6-v2", use_openai: bool = False):
self.use_openai = use_openai
self.model_name = model_name
if self.use_openai:
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
raise ValueError("OPENAI_API_KEY not found in environment.")
self.client = OpenAI(api_key=api_key)
else:
print(f"Loading local embedding model: {model_name}")
# Force CPU to avoid ZeroGPU conflicts during ingestion
self.model = SentenceTransformer(model_name, device="cpu")
def embed(self, texts: List[str]) -> np.ndarray:
if not texts:
return np.array([])
if self.use_openai:
# Batching might be needed for very large lists, but keeping simple for now
response = self.client.embeddings.create(input=texts, model=self.model_name)
embeddings = [data.embedding for data in response.data]
return np.array(embeddings, dtype='float32')
else:
return self.model.encode(texts, convert_to_numpy=True)
def get_embedder():
# Factory to get configured embedder
# Prefer OpenAI if specified, else local
use_openai = os.getenv("USE_OPENAI_EMBEDDINGS", "false").lower() == "true"
model_name = "text-embedding-3-small" if use_openai else "all-MiniLM-L6-v2"
return Embedder(model_name=model_name, use_openai=use_openai)
|