import numpy as np
from sentence_transformers import SentenceTransformer

# 1) 정답 — 한 문장 임베딩
model = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")
emb = model.encode("인공지능이 세상을 바꾸고 있다.")
print(f"벡터 차원: {emb.shape[0]}")          # 예: 384
print(f"처음 10개 값: {emb[:10].round(4)}")  # 실수들의 배열

# 2) 정답 — 여러 문서를 한 번에 임베딩
documents = [
    "Python은 강력한 프로그래밍 언어입니다.",
    "OpenAI는 AI 연구를 선도하는 기업입니다.",
    "FAISS는 벡터 검색을 위한 라이브러리입니다.",
]
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
doc_embeddings = np.array(embedding_model.encode(documents))
print(f"doc_embeddings shape: {doc_embeddings.shape}")  # (3, 384)

# 3) 정답 — 더 정확한 모델로 교체 (참고)
# all-MiniLM-L6-v2  : 384차원,  80MB, 빠름 / 품질 보통
# all-mpnet-base-v2 : 768차원, 420MB, 느림 / 품질 우수
# (모델을 바꾸면 '문서도 새 모델로 다시 임베딩'해야 한다 → 같은 공간 규칙)
model_name = "all-mpnet-base-v2"
embedding_model = SentenceTransformer(model_name)
doc_embeddings = np.array(embedding_model.encode(documents), dtype=np.float32)


# 4) 정답 — 문서와 질문은 '같은 로컬 모델'로 (필수 규칙)
def rag_query(user_query):
    # 질문도 같은 로컬 모델로 임베딩 → 코사인 유사도가 가장 큰 문서 (faiss 없이 numpy로)
    query_vec = np.array(embedding_model.encode(user_query), dtype=np.float32)
    sims = doc_embeddings @ query_vec / (
        np.linalg.norm(doc_embeddings, axis=1) * np.linalg.norm(query_vec))
    return documents[int(np.argmax(sims))]


if __name__ == "__main__":
    q = "벡터 검색용 라이브러리가 뭐야?"
    print("\n[로컬 모델 검색 데모]")
    print("질문:", q)
    print("검색된 문서:", rag_query(q))
