peekabook-api / app /embedding /embedder.py
lael
feat: update to graph_main with HyDE RAG, user_id isolation, GPU auto-detect
7d2a564
Raw
History Blame Contribute Delete
3.27 kB
# src/embedding/embedder.py
from abc import ABC, abstractmethod
from dotenv import load_dotenv
import os
class BaseEmbedder(ABC):
@abstractmethod
def embed(self, text: str) -> list[float]:
pass
@abstractmethod
def embed_batch(self, texts: list[str]) -> list[list[float]]:
pass
# ─────────────────────────────────────────
# 둜컬 λͺ¨λΈ
# ─────────────────────────────────────────
class LocalEmbedder(BaseEmbedder):
"""
둜컬 λͺ¨λΈ μž„λ² λ”© (sentence-transformers 기반)
Parameters
----------
model_name : HuggingFace λͺ¨λΈλͺ… (κΈ°λ³Έ: BAAI/bge-m3, vector_size=1024)
예)
embedder = LocalEmbedder()
embedder = LocalEmbedder("BAAI/bge-m3")
"""
def __init__(self, model_name="BAAI/bge-m3"):
import torch
from sentence_transformers import SentenceTransformer
device = "cuda" if torch.cuda.is_available() else "cpu"
self.model = SentenceTransformer(model_name, device=device)
def embed(self, text: str) -> list[float]:
return self.model.encode(text).tolist() # 좜λ ₯ ν˜•μ‹ : 단일 리슀트
def embed_batch(self, texts: list[str]) -> list[list[float]]:
return self.model.encode(texts).tolist() # 좜λ ₯ ν˜•μ‹ : 벑터 μ—¬λŸ¬ 개λ₯Ό 담은 리슀트
# ─────────────────────────────────────────
# API 기반 λͺ¨λΈ
# ─────────────────────────────────────────
class APIEmbedder(BaseEmbedder, ABC):
"""
API 기반 μž„λ² λ”© 베이슀 클래슀.
μ™ΈλΆ€ APIλ₯Ό μ‚¬μš©ν•˜λŠ” μž„λ² λ”λŠ” 이 클래슀λ₯Ό 상속해 κ΅¬ν˜„.
κ΅¬ν˜„μ²΄ λͺ©λ‘:
- OpenAIEmbedder : OpenAI Embeddings API
"""
pass
# class OpenAIEmbedder(APIEmbedder):
# """
# OpenAI Embeddings API
# Parameters
# ----------
# model_name : μ‚¬μš©ν•  λͺ¨λΈ (κΈ°λ³Έ: text-embedding-3-small)
# dimensions : 좜λ ₯ 벑터 차원 수 (κΈ°λ³Έ: 1024 β†’ bge-m3κ³Ό λ™μΌν•˜κ²Œ 맞좰 μ»¬λ ‰μ…˜ μž¬μ‚¬μš© κ°€λŠ₯)
# 예)
# embedder = OpenAIEmbedder()
# embedder = OpenAIEmbedder("text-embedding-3-large", dimensions=1024)
# """
# def __init__(self, model_name="text-embedding-3-small", dimensions=1024):
# from openai import OpenAI
# load_dotenv()
# self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
# self.model_name = model_name
# self.dimensions = dimensions
# def embed(self, text: str) -> list[float]:
# response = self.client.embeddings.create(
# model=self.model_name,
# input=text,
# dimensions=self.dimensions
# )
# return response.data[0].embedding
# def embed_batch(self, texts: list[str]) -> list[list[float]]:
# response = self.client.embeddings.create(
# model=self.model_name,
# input=texts,
# dimensions=self.dimensions
# )
# return [item.embedding for item in response.data]