임베딩 모델 완전 가이드: 벡터 검색·RAG·Sentence Transformers 실전 활용

# 직관적 이해: 단어를 벡터로 표현
# "왕" = [0.2, 0.8, 0.1, ...]
# "여왕" = [0.3, 0.9, 0.1, ...]
# "남자" = [0.1, 0.2, 0.8, ...]

# 유명한 벡터 산술: king - man + woman ≈ queen
import numpy as np

king = np.array([0.2, 0.8, 0.1, 0.5])
man = np.array([0.1, 0.2, 0.8, 0.4])
woman = np.array([0.15, 0.25, 0.85, 0.6])
queen = np.array([0.3, 0.9, 0.1, 0.7])

result = king - man + woman
print(f"king - man + woman = {result}")
print(f"queen              = {queen}")
# 두 벡터가 매우 유사함을 확인할 수 있다

# 모델 선택 의사결정 트리
def select_embedding_model(requirements):
    """요구사항에 따른 임베딩 모델 선택 가이드"""

    if requirements.get("budget") == "unlimited":
        if requirements.get("max_performance"):
            return "GTE-Qwen2-7B-instruct (자체 호스팅) 또는 Voyage-3 (API)"
        return "text-embedding-3-large (OpenAI API)"

    if requirements.get("multilingual"):
        if requirements.get("self_hosted"):
            return "BGE-M3 (Dense+Sparse 하이브리드)"
        return "Cohere embed-v3 multilingual"

    if requirements.get("low_latency"):
        if requirements.get("self_hosted"):
            return "nomic-embed-text-v1.5 (경량 137M)"
        return "text-embedding-3-small (OpenAI API)"

    if requirements.get("domain_specific"):
        return "Sentence Transformers + 파인튜닝 (기본 모델: BGE or E5)"

    # 기본 추천
    return "text-embedding-3-small (비용 효율적 범용 선택)"

from sentence_transformers import SentenceTransformer
import numpy as np

# 모델 로드
model = SentenceTransformer('BAAI/bge-large-en-v1.5')

# 단일 문장 임베딩
sentence = "Embedding models convert text into numerical vectors."
embedding = model.encode(sentence)
print(f"차원: {embedding.shape}")  # (1024,)

# 배치 임베딩
sentences = [
    "임베딩 모델은 텍스트를 수치 벡터로 변환한다.",
    "벡터 검색은 유사한 문서를 빠르게 찾아준다.",
    "RAG는 검색 기반 생성 기법이다.",
    "오늘 날씨가 매우 좋다.",
]

embeddings = model.encode(sentences, batch_size=32, show_progress_bar=True)
print(f"임베딩 행렬 크기: {embeddings.shape}")  # (4, 1024)

# 유사도 계산
from sentence_transformers.util import cos_sim

similarity_matrix = cos_sim(embeddings, embeddings)
print("유사도 행렬:")
print(similarity_matrix.numpy().round(3))

from sentence_transformers import SentenceTransformer

# BGE-M3: 100개 이상의 언어를 지원하는 다국어 임베딩 모델
model = SentenceTransformer('BAAI/bge-m3')

# 다국어 문장 임베딩
sentences = [
    "Machine learning is transforming the world.",        # 영어
    "머신러닝이 세상을 변화시키고 있다.",                      # 한국어
    "機械学習が世界を変えている。",                           # 일본어
    "机器学习正在改变世界。",                                # 중국어
]

embeddings = model.encode(sentences, normalize_embeddings=True)

# 다국어 간 유사도 확인
from sentence_transformers.util import cos_sim

similarities = cos_sim(embeddings, embeddings)
print("다국어 유사도 행렬:")
for i, s1 in enumerate(sentences):
    for j, s2 in enumerate(sentences):
        if i < j:
            print(f"  '{s1[:20]}...' <-> '{s2[:20]}...': {similarities[i][j]:.4f}")
# 같은 의미의 다른 언어 문장들이 높은 유사도를 보인다

from openai import OpenAI
import numpy as np

client = OpenAI()

def get_openai_embeddings(texts, model="text-embedding-3-small", dimensions=None):
    """OpenAI 임베딩 생성 (Matryoshka 차원 축소 지원)"""
    kwargs = {"input": texts, "model": model}
    if dimensions:
        kwargs["dimensions"] = dimensions

    response = client.embeddings.create(**kwargs)
    return np.array([item.embedding for item in response.data])

# 기본 사용
texts = ["임베딩 모델의 원리", "벡터 검색 시스템 구축"]
embeddings_full = get_openai_embeddings(texts, model="text-embedding-3-large")
print(f"전체 차원: {embeddings_full.shape}")  # (2, 3072)

# Matryoshka: 차원 축소로 비용/속도 최적화
embeddings_256 = get_openai_embeddings(
    texts, model="text-embedding-3-large", dimensions=256
)
print(f"축소 차원: {embeddings_256.shape}")  # (2, 256)

# 코사인 유사도 비교
def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

sim_full = cosine_similarity(embeddings_full[0], embeddings_full[1])
sim_256 = cosine_similarity(embeddings_256[0], embeddings_256[1])
print(f"전체 차원 유사도: {sim_full:.4f}")
print(f"256차원 유사도: {sim_256:.4f}")

import faiss
import numpy as np
import time

# 테스트 데이터 생성
np.random.seed(42)
dimension = 1024
num_vectors = 1_000_000
num_queries = 100

# 정규화된 랜덤 벡터 생성
data = np.random.randn(num_vectors, dimension).astype('float32')
faiss.normalize_L2(data)
queries = np.random.randn(num_queries, dimension).astype('float32')
faiss.normalize_L2(queries)

# 1. Flat Index (정확 검색, Brute-force)
print("=== Flat Index (Exact Search) ===")
index_flat = faiss.IndexFlatIP(dimension)  # Inner Product (코사인 유사도)
index_flat.add(data)

start = time.time()
D_exact, I_exact = index_flat.search(queries, 10)
flat_time = time.time() - start
print(f"검색 시간: {flat_time:.3f}초")
print(f"Recall@10: 1.000 (정확 검색)")

# 2. IVF (Inverted File Index)
print("\n=== IVF Index ===")
nlist = 1024  # 클러스터 수
quantizer = faiss.IndexFlatIP(dimension)
index_ivf = faiss.IndexIVFFlat(quantizer, dimension, nlist, faiss.METRIC_INNER_PRODUCT)
index_ivf.train(data)
index_ivf.add(data)
index_ivf.nprobe = 32  # 검색할 클러스터 수

start = time.time()
D_ivf, I_ivf = index_ivf.search(queries, 10)
ivf_time = time.time() - start

# Recall 계산
recall = np.mean([len(set(I_ivf[i]) & set(I_exact[i])) / 10 for i in range(num_queries)])
print(f"검색 시간: {ivf_time:.3f}초 (x{flat_time/ivf_time:.1f} 빠름)")
print(f"Recall@10: {recall:.3f}")

# 3. HNSW (Hierarchical Navigable Small World)
print("\n=== HNSW Index ===")
index_hnsw = faiss.IndexHNSWFlat(dimension, 32)  # M=32
index_hnsw.hnsw.efConstruction = 200
index_hnsw.hnsw.efSearch = 64
index_hnsw.metric_type = faiss.METRIC_INNER_PRODUCT
index_hnsw.add(data)

start = time.time()
D_hnsw, I_hnsw = index_hnsw.search(queries, 10)
hnsw_time = time.time() - start

recall = np.mean([len(set(I_hnsw[i]) & set(I_exact[i])) / 10 for i in range(num_queries)])
print(f"검색 시간: {hnsw_time:.3f}초 (x{flat_time/hnsw_time:.1f} 빠름)")
print(f"Recall@10: {recall:.3f}")

# 4. IVF-PQ (Product Quantization)
print("\n=== IVF-PQ Index (메모리 최적화) ===")
m = 64  # 서브벡터 수
nbits = 8  # 코드북 비트 수
index_ivfpq = faiss.IndexIVFPQ(quantizer, dimension, nlist, m, nbits)
index_ivfpq.train(data)
index_ivfpq.add(data)
index_ivfpq.nprobe = 32

start = time.time()
D_pq, I_pq = index_ivfpq.search(queries, 10)
pq_time = time.time() - start

recall = np.mean([len(set(I_pq[i]) & set(I_exact[i])) / 10 for i in range(num_queries)])
print(f"검색 시간: {pq_time:.3f}초 (x{flat_time/pq_time:.1f} 빠름)")
print(f"Recall@10: {recall:.3f}")
print(f"메모리: Flat={data.nbytes/1e9:.1f}GB, PQ={index_ivfpq.sa_code_size()*num_vectors/1e9:.3f}GB")

import chromadb
from chromadb.utils import embedding_functions

# Chroma 클라이언트 생성
client = chromadb.PersistentClient(path="./chroma_db")

# Sentence Transformers 임베딩 함수 설정
embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="BAAI/bge-m3"
)

# 컬렉션 생성
collection = client.get_or_create_collection(
    name="tech_documents",
    embedding_function=embedding_fn,
    metadata={"hnsw:space": "cosine"}  # 코사인 유사도 사용
)

# 문서 추가
documents = [
    "Kubernetes는 컨테이너 오케스트레이션 플랫폼이다.",
    "Docker는 애플리케이션을 컨테이너로 패키징하는 도구다.",
    "Prometheus는 메트릭 기반의 모니터링 시스템이다.",
    "Grafana는 데이터 시각화 및 대시보드 도구다.",
    "Terraform은 인프라를 코드로 관리하는 IaC 도구다.",
    "임베딩 모델은 텍스트를 벡터로 변환한다.",
    "RAG는 검색 증강 생성 기법으로 LLM의 환각을 줄인다.",
]

collection.add(
    documents=documents,
    ids=[f"doc_{i}" for i in range(len(documents))],
    metadatas=[
        {"category": "kubernetes"},
        {"category": "docker"},
        {"category": "monitoring"},
        {"category": "monitoring"},
        {"category": "iac"},
        {"category": "ai"},
        {"category": "ai"},
    ]
)

# 시맨틱 검색
results = collection.query(
    query_texts=["컨테이너 관련 기술은?"],
    n_results=3
)
print("검색 결과:")
for doc, score in zip(results["documents"][0], results["distances"][0]):
    print(f"  [{score:.4f}] {doc}")

# 메타데이터 필터링 + 시맨틱 검색
results_filtered = collection.query(
    query_texts=["모니터링 도구"],
    n_results=2,
    where={"category": "monitoring"}
)
print("\n필터링된 검색 결과:")
for doc in results_filtered["documents"][0]:
    print(f"  {doc}")

import numpy as np

def cosine_similarity(a, b):
    """코사인 유사도: 벡터 방향의 유사성 측정"""
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

def dot_product(a, b):
    """내적: 정규화된 벡터에서는 코사인 유사도와 동일"""
    return np.dot(a, b)

def euclidean_distance(a, b):
    """유클리드 거리: 벡터 간 직선 거리"""
    return np.linalg.norm(a - b)

def manhattan_distance(a, b):
    """맨하탄 거리: 차원별 절대 차이의 합"""
    return np.sum(np.abs(a - b))

# 예시 벡터
a = np.array([1.0, 2.0, 3.0])
b = np.array([1.0, 2.0, 3.1])
c = np.array([-1.0, -2.0, -3.0])

print("벡터 a와 b (매우 유사):")
print(f"  코사인 유사도:  {cosine_similarity(a, b):.4f}")
print(f"  유클리드 거리:  {euclidean_distance(a, b):.4f}")
print(f"  내적:          {dot_product(a, b):.4f}")

print("\n벡터 a와 c (반대 방향):")
print(f"  코사인 유사도:  {cosine_similarity(a, c):.4f}")
print(f"  유클리드 거리:  {euclidean_distance(a, c):.4f}")
print(f"  내적:          {dot_product(a, c):.4f}")

from sentence_transformers import SentenceTransformer, util
import torch

class SemanticSearchEngine:
    def __init__(self, model_name="BAAI/bge-m3"):
        self.model = SentenceTransformer(model_name)
        self.documents = []
        self.embeddings = None

    def index_documents(self, documents):
        """문서를 인덱싱하여 임베딩 생성"""
        self.documents = documents
        self.embeddings = self.model.encode(
            documents,
            convert_to_tensor=True,
            normalize_embeddings=True,
            show_progress_bar=True
        )
        print(f"{len(documents)}개 문서 인덱싱 완료 (차원: {self.embeddings.shape[1]})")

    def search(self, query, top_k=5):
        """시맨틱 검색 수행"""
        query_embedding = self.model.encode(
            query,
            convert_to_tensor=True,
            normalize_embeddings=True
        )

        # 코사인 유사도 계산
        scores = util.cos_sim(query_embedding, self.embeddings)[0]

        # 상위 k개 결과 반환
        top_results = torch.topk(scores, k=min(top_k, len(self.documents)))

        results = []
        for score, idx in zip(top_results.values, top_results.indices):
            results.append({
                "document": self.documents[idx],
                "score": score.item(),
                "index": idx.item()
            })
        return results

    def search_with_reranking(self, query, top_k=5, initial_k=20):
        """2단계 검색: 임베딩 검색 + 리랭킹"""
        from sentence_transformers import CrossEncoder

        # 1단계: 임베딩 기반 후보 검색
        candidates = self.search(query, top_k=initial_k)

        # 2단계: Cross-encoder로 리랭킹
        reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')
        pairs = [(query, c["document"]) for c in candidates]
        rerank_scores = reranker.predict(pairs)

        # 리랭킹 결과 반환
        for i, score in enumerate(rerank_scores):
            candidates[i]["rerank_score"] = float(score)

        reranked = sorted(candidates, key=lambda x: x["rerank_score"], reverse=True)
        return reranked[:top_k]

# 사용 예시
engine = SemanticSearchEngine()

documents = [
    "Python은 데이터 과학과 머신러닝에 가장 많이 사용되는 프로그래밍 언어다.",
    "JavaScript는 웹 개발의 핵심 언어로 Node.js를 통해 서버 사이드에서도 사용된다.",
    "Kubernetes는 컨테이너화된 애플리케이션의 배포, 스케일링, 관리를 자동화한다.",
    "PostgreSQL은 강력한 오픈소스 관계형 데이터베이스 관리 시스템이다.",
    "TensorFlow와 PyTorch는 딥러닝 모델 개발에 가장 널리 사용되는 프레임워크다.",
    "Redis는 인메모리 데이터 구조 저장소로 캐싱과 메시지 브로커로 활용된다.",
    "Docker는 애플리케이션과 종속성을 컨테이너로 패키징하여 이식성을 제공한다.",
    "GraphQL은 REST의 대안으로 클라이언트가 필요한 데이터만 요청할 수 있게 한다.",
]

engine.index_documents(documents)

# 시맨틱 검색
query = "딥러닝 개발에 어떤 도구를 사용해야 하나요?"
results = engine.search(query, top_k=3)
print(f"\n쿼리: {query}")
for r in results:
    print(f"  [{r['score']:.4f}] {r['document']}")

from sentence_transformers import SentenceTransformer, CrossEncoder
from openai import OpenAI
import chromadb
from chromadb.utils import embedding_functions
import tiktoken
from typing import List, Dict

class RAGPipeline:
    def __init__(
        self,
        embedding_model: str = "BAAI/bge-m3",
        reranker_model: str = "cross-encoder/ms-marco-MiniLM-L-12-v2",
        llm_model: str = "gpt-4o",
    ):
        self.embedder = SentenceTransformer(embedding_model)
        self.reranker = CrossEncoder(reranker_model)
        self.llm_client = OpenAI()
        self.llm_model = llm_model

        # Chroma 벡터 DB 초기화
        self.chroma_client = chromadb.PersistentClient(path="./rag_db")
        self.collection = self.chroma_client.get_or_create_collection(
            name="rag_documents",
            metadata={"hnsw:space": "cosine"}
        )

    def chunk_text(self, text: str, chunk_size: int = 512, overlap: int = 50) -> List[str]:
        """텍스트를 토큰 기반으로 청크 분할"""
        tokenizer = tiktoken.get_encoding("cl100k_base")
        tokens = tokenizer.encode(text)
        chunks = []

        start = 0
        while start < len(tokens):
            end = start + chunk_size
            chunk_tokens = tokens[start:end]
            chunk_text = tokenizer.decode(chunk_tokens)
            chunks.append(chunk_text)
            start = end - overlap  # 오버랩 적용

        return chunks

    def ingest_documents(self, documents: List[Dict[str, str]]):
        """문서를 청크 분할하고 벡터 DB에 저장"""
        all_chunks = []
        all_ids = []
        all_metadatas = []

        for doc_idx, doc in enumerate(documents):
            chunks = self.chunk_text(doc["content"])
            for chunk_idx, chunk in enumerate(chunks):
                all_chunks.append(chunk)
                all_ids.append(f"doc{doc_idx}_chunk{chunk_idx}")
                all_metadatas.append({
                    "source": doc.get("source", "unknown"),
                    "doc_index": doc_idx,
                    "chunk_index": chunk_idx,
                })

        # 임베딩 생성 및 저장
        embeddings = self.embedder.encode(all_chunks, normalize_embeddings=True)

        self.collection.add(
            documents=all_chunks,
            embeddings=embeddings.tolist(),
            ids=all_ids,
            metadatas=all_metadatas,
        )
        print(f"{len(documents)}개 문서 -> {len(all_chunks)}개 청크 인덱싱 완료")

    def retrieve(self, query: str, top_k: int = 10) -> List[Dict]:
        """벡터 유사도 기반 검색"""
        query_embedding = self.embedder.encode(
            [query], normalize_embeddings=True
        ).tolist()

        results = self.collection.query(
            query_embeddings=query_embedding,
            n_results=top_k,
        )

        retrieved = []
        for i in range(len(results["documents"][0])):
            retrieved.append({
                "text": results["documents"][0][i],
                "metadata": results["metadatas"][0][i],
                "distance": results["distances"][0][i],
            })
        return retrieved

    def rerank(self, query: str, candidates: List[Dict], top_k: int = 5) -> List[Dict]:
        """Cross-encoder 기반 리랭킹"""
        pairs = [(query, c["text"]) for c in candidates]
        scores = self.reranker.predict(pairs)

        for i, score in enumerate(scores):
            candidates[i]["rerank_score"] = float(score)

        reranked = sorted(candidates, key=lambda x: x["rerank_score"], reverse=True)
        return reranked[:top_k]

    def generate(self, query: str, context_docs: List[Dict]) -> str:
        """검색된 컨텍스트를 기반으로 LLM 응답 생성"""
        context = "\n\n---\n\n".join([doc["text"] for doc in context_docs])

        messages = [
            {
                "role": "system",
                "content": (
                    "You are a helpful assistant. Answer the question based on "
                    "the provided context. If the context doesn't contain "
                    "relevant information, say so."
                ),
            },
            {
                "role": "user",
                "content": f"Context:\n{context}\n\nQuestion: {query}",
            },
        ]

        response = self.llm_client.chat.completions.create(
            model=self.llm_model,
            messages=messages,
            temperature=0.1,
        )
        return response.choices[0].message.content

    def query(self, question: str, top_k_retrieve: int = 10, top_k_rerank: int = 5) -> str:
        """전체 RAG 파이프라인 실행"""
        # 1. 검색
        candidates = self.retrieve(question, top_k=top_k_retrieve)
        print(f"1단계 검색: {len(candidates)}개 후보")

        # 2. 리랭킹
        reranked = self.rerank(question, candidates, top_k=top_k_rerank)
        print(f"2단계 리랭킹: 상위 {len(reranked)}개 선택")

        # 3. 생성
        answer = self.generate(question, reranked)
        return answer

# 사용 예시
rag = RAGPipeline()

# 문서 인제스트
documents = [
    {"content": "긴 기술 문서 내용...", "source": "tech_doc_1.pdf"},
    {"content": "또 다른 문서 내용...", "source": "tech_doc_2.pdf"},
]
rag.ingest_documents(documents)

# 질의
answer = rag.query("임베딩 모델의 차원 수는 성능에 어떤 영향을 미치나요?")
print(f"\n답변: {answer}")

from rank_bm25 import BM25Okapi
import numpy as np

class HybridSearchEngine:
    """Dense(임베딩) + Sparse(BM25) 하이브리드 검색"""

    def __init__(self, embedding_model="BAAI/bge-m3"):
        self.embedder = SentenceTransformer(embedding_model)
        self.documents = []
        self.embeddings = None
        self.bm25 = None

    def index(self, documents):
        self.documents = documents

        # Dense: 임베딩 생성
        self.embeddings = self.embedder.encode(
            documents, normalize_embeddings=True, convert_to_tensor=True
        )

        # Sparse: BM25 인덱스 구축
        tokenized = [doc.split() for doc in documents]
        self.bm25 = BM25Okapi(tokenized)

    def search(self, query, top_k=5, alpha=0.7):
        """하이브리드 검색 (alpha: dense 가중치, 1-alpha: sparse 가중치)"""
        # Dense 검색
        query_emb = self.embedder.encode(
            query, normalize_embeddings=True, convert_to_tensor=True
        )
        dense_scores = util.cos_sim(query_emb, self.embeddings)[0].cpu().numpy()

        # Sparse 검색 (BM25)
        sparse_scores = self.bm25.get_scores(query.split())

        # 정규화
        if dense_scores.max() > 0:
            dense_scores = dense_scores / dense_scores.max()
        if sparse_scores.max() > 0:
            sparse_scores = sparse_scores / sparse_scores.max()

        # 가중 합산
        hybrid_scores = alpha * dense_scores + (1 - alpha) * sparse_scores

        # 상위 k개 반환
        top_indices = np.argsort(hybrid_scores)[::-1][:top_k]
        return [
            {
                "document": self.documents[i],
                "score": hybrid_scores[i],
                "dense_score": dense_scores[i],
                "sparse_score": sparse_scores[i],
            }
            for i in top_indices
        ]

from sentence_transformers import (
    SentenceTransformer,
    InputExample,
    losses,
    evaluation,
)
from torch.utils.data import DataLoader

# 기본 모델 로드
model = SentenceTransformer('BAAI/bge-base-en-v1.5')

# 학습 데이터 준비 (anchor, positive, negative)
train_examples = [
    # (쿼리, 관련 문서, 비관련 문서)
    InputExample(texts=[
        "How to deploy a Kubernetes pod?",
        "kubectl apply -f pod.yaml creates a new pod in the cluster.",
        "Python is a popular programming language for data science."
    ]),
    InputExample(texts=[
        "What is a Docker container?",
        "A Docker container is a lightweight, standalone executable package.",
        "Machine learning models require large datasets for training."
    ]),
    InputExample(texts=[
        "How does Redis caching work?",
        "Redis stores data in memory for fast read/write access as a cache layer.",
        "Kubernetes orchestrates containerized applications across clusters."
    ]),
    # ... 수천 ~ 수만 개의 학습 예시
]

# DataLoader 생성
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)

# 손실 함수: TripletLoss (anchor, positive, negative)
train_loss = losses.TripletLoss(model=model)

# 평가 데이터
eval_examples = [
    InputExample(texts=["query1", "relevant_doc1"], label=1.0),
    InputExample(texts=["query2", "irrelevant_doc2"], label=0.0),
]
evaluator = evaluation.EmbeddingSimilarityEvaluator.from_input_examples(
    eval_examples, name="domain-eval"
)

# 파인튜닝 실행
model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    evaluator=evaluator,
    epochs=3,
    warmup_steps=100,
    evaluation_steps=500,
    output_path="./finetuned_embedding_model",
    save_best_model=True,
)

print("파인튜닝 완료!")

# 파인튜닝된 모델 로드 및 사용
finetuned_model = SentenceTransformer('./finetuned_embedding_model')
embeddings = finetuned_model.encode(["도메인 특화 쿼리"])

from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader

model = SentenceTransformer('BAAI/bge-base-en-v1.5')

# (query, positive_passage) 쌍만으로 학습 가능
# in-batch negatives를 자동으로 활용
train_examples = [
    InputExample(texts=["What is embedding?", "An embedding is a vector representation of data."]),
    InputExample(texts=["How does HNSW work?", "HNSW builds a hierarchical graph for approximate nearest neighbor search."]),
    InputExample(texts=["What is RAG?", "RAG retrieves relevant documents and uses them to augment LLM generation."]),
    # ... 더 많은 (query, positive) 쌍
]

train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=32)

# MultipleNegativesRankingLoss: 배치 내 다른 positive를 negative로 활용
train_loss = losses.MultipleNegativesRankingLoss(model=model)

model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=3,
    warmup_steps=100,
    output_path="./mnrl_finetuned_model",
)

from sentence_transformers import SentenceTransformer
import numpy as np

# Matryoshka Representation Learning (MRL) 지원 모델
model = SentenceTransformer('nomic-ai/nomic-embed-text-v1.5')

texts = [
    "Vector databases store embeddings for similarity search.",
    "Embedding models convert text into numerical representations.",
    "RAG systems combine retrieval with language generation.",
]

# 전체 차원 임베딩
full_embeddings = model.encode(texts)
print(f"전체 차원: {full_embeddings.shape[1]}")  # 768

# Matryoshka: 원하는 차원으로 절삭 후 정규화
def truncate_embeddings(embeddings, target_dim):
    """Matryoshka 방식으로 차원 축소"""
    truncated = embeddings[:, :target_dim]
    # L2 정규화
    norms = np.linalg.norm(truncated, axis=1, keepdims=True)
    return truncated / norms

# 다양한 차원에서의 유사도 비교
for dim in [64, 128, 256, 512, 768]:
    reduced = truncate_embeddings(full_embeddings, dim)
    sim = np.dot(reduced[0], reduced[1])  # 정규화된 벡터의 내적 = 코사인 유사도
    print(f"  차원 {dim:>4}: 유사도 = {sim:.4f}")

import numpy as np

def scalar_quantize_int8(embeddings):
    """스칼라 양자화: float32 -> int8 (메모리 75% 절감)"""
    min_val = embeddings.min(axis=0)
    max_val = embeddings.max(axis=0)
    scale = (max_val - min_val) / 255.0

    quantized = ((embeddings - min_val) / scale).astype(np.int8)
    return quantized, min_val, scale

def scalar_dequantize_int8(quantized, min_val, scale):
    """역양자화: int8 -> float32"""
    return quantized.astype(np.float32) * scale + min_val

def binary_quantize(embeddings):
    """이진 양자화: float32 -> 1bit (메모리 32배 절감)"""
    return (embeddings > 0).astype(np.uint8)

# 메모리 비교
num_vectors = 1_000_000
dimension = 1024
embeddings = np.random.randn(num_vectors, dimension).astype(np.float32)

print(f"원본 (float32): {embeddings.nbytes / 1e9:.2f} GB")

quantized, _, _ = scalar_quantize_int8(embeddings)
print(f"int8 양자화:    {quantized.nbytes / 1e9:.2f} GB")

binary = binary_quantize(embeddings)
print(f"이진 양자화:    {binary.nbytes / 1e9:.2f} GB")
# 원본 (float32): 4.10 GB
# int8 양자화:    1.02 GB
# 이진 양자화:    1.02 GB (실제 비트 패킹 시 0.13 GB)

import asyncio
from sentence_transformers import SentenceTransformer
from functools import lru_cache
import hashlib

class OptimizedEmbeddingService:
    def __init__(self, model_name="BAAI/bge-m3", cache_size=10000):
        self.model = SentenceTransformer(model_name)
        self.cache = {}
        self.cache_size = cache_size

    def _get_cache_key(self, text):
        return hashlib.md5(text.encode()).hexdigest()

    def encode_with_cache(self, texts, batch_size=64):
        """캐싱을 적용한 임베딩 생성"""
        uncached_texts = []
        uncached_indices = []
        results = [None] * len(texts)

        # 캐시 히트 확인
        for i, text in enumerate(texts):
            key = self._get_cache_key(text)
            if key in self.cache:
                results[i] = self.cache[key]
            else:
                uncached_texts.append(text)
                uncached_indices.append(i)

        # 캐시 미스 텍스트만 배치 임베딩
        if uncached_texts:
            new_embeddings = self.model.encode(
                uncached_texts,
                batch_size=batch_size,
                normalize_embeddings=True,
            )

            for idx, emb in zip(uncached_indices, new_embeddings):
                key = self._get_cache_key(texts[idx])
                self.cache[key] = emb
                results[idx] = emb

                # 캐시 크기 관리
                if len(self.cache) > self.cache_size:
                    oldest_key = next(iter(self.cache))
                    del self.cache[oldest_key]

        return results

    def get_cache_stats(self):
        return {"cache_size": len(self.cache), "max_size": self.cache_size}