Spaces:
Running
Running
| # src/embedding/embedder.py | |
| from abc import ABC, abstractmethod | |
| from dotenv import load_dotenv | |
| import os | |
| class BaseEmbedder(ABC): | |
| def embed(self, text: str) -> list[float]: | |
| pass | |
| 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] | |