실습 3 — 벡터DB 에 넣고 찾기

import numpy as np
from sentence_transformers import SentenceTransformer

text = open("it_helpdesk.md", encoding="utf-8").read()
chunks = ["## " + s.strip() for s in text.split("\n## ")[1:]]
emb = SentenceTransformer("intfloat/multilingual-e5-small", device="cpu")
E = emb.encode(["passage: " + c for c in chunks], normalize_embeddings=True)
def title(c):
    return c.split("\n")[0].removeprefix("## ")
question = "핸드폰을 새로 샀는데 인증은 어떻게 옮기나요?"
q = emb.encode(["query: " + question], normalize_embeddings=True)     # 모양 (1, 384): FAISS 는 2차원 배열을 받는다
print(E.dtype, E.shape, q.shape)

import faiss

index = faiss.IndexFlatIP(E.shape[1])   # 384차원, 내적(IP)으로 모두 견주는 인덱스
index.add(E)                            # 조각 30개를 0~29 번으로 넣는다
faiss.write_index(index, "out_helpdesk.faiss")

loaded = faiss.read_index("out_helpdesk.faiss")    # 파일에서 다시 불러온다
D, I = loaded.search(q, 3)              # D: 점수, I: 조각 번호. 둘 다 (질문 수, k) 모양
print("넣은 벡터 수:", loaded.ntotal)
for s, i in zip(D[0], I[0]):
    print(round(float(s), 3), i, title(chunks[i]))    # 번호에 맞는 원문은 우리가 든 chunks 에서 찾는다

import chromadb

client = chromadb.PersistentClient(path="out_chroma")    # 이 폴더에 저장한다
# 같은 이름의 컬렉션이 있으면 지우고 새로 만든다(다시 실행해도 겹치지 않게)
if "helpdesk" in [c.name for c in client.list_collections()]:
    client.delete_collection("helpdesk")
col = client.create_collection("helpdesk", metadata={"hnsw:space": "cosine"})   # 거리: 코사인
col.add(ids=[f"s{i}" for i in range(len(chunks))],
        documents=chunks,                                  # 원문도 함께 넣는다
        embeddings=E.tolist(),                             # 우리가 만든 임베딩을 직접 넣는다
        metadatas=[{"제목": title(c)} for c in chunks])
print("컬렉션의 조각 수:", col.count())

r = col.query(query_embeddings=q.tolist(), n_results=3)
for i, d, m in zip(r["ids"][0], r["distances"][0], r["metadatas"][0]):
    print(i, round(d, 3), m["제목"])                       # distance = 1 - 코사인 유사도

# 제목이 '와이파이' 인 조각 안에서만 찾는다
r = col.query(query_embeddings=q.tolist(), n_results=1, where={"제목": "와이파이"})
print(r["ids"][0], r["documents"][0][0][:40])

# 다시 열어도 남아 있는지: 새 클라이언트로 같은 폴더를 연다
again = chromadb.PersistentClient(path="out_chroma").get_collection("helpdesk")
print("다시 연 컬렉션의 조각 수:", again.count())