환경

구성	버전
OS	Ubuntu 22.04
Python	3.10.12

데이터 구성

소설, 웹툰 데이터로

제목, 작가, 소개글, 장르, 태그, 썸네일, url, 조회수, 추천수, 평점, 연재수, 완결 여부, 나이, 출판사, 플랫폼, 타입이 있다.

import chromadb
import pandas as pd
from tqdm import tqdm

# Chroma DB 인스턴스 생성
chroma_client = chromadb.PersistentClient(path="./chroma_db")

# 컬렉션 생성 (벡터를 저장할 공간)
collection = chroma_client.get_or_create_collection(name="work_collection")

print('-- 1. 데이터 불러오기 --')
df = pd.read_csv('datas/train_data.csv', encoding='utf-8')
print(df.head())

# 1. 저장할 자료 준비
documents = [
    f"제목: {row['title']}, 작가: {row['author']}, 소개글: {row['summary']}, 장르: {row['genre']}, 태그: {row['tags']}"
        for _, row in df.iterrows()
]


print('\n-- 2. 데이터 변환 및 저장 준비 --')
# 배치 크기 설정 (한 번에 처리할 데이터 수)
batch_size = 100

# 전체 데이터를 배치로 나누기
total_batches = (len(df) + batch_size - 1) // batch_size
print(f"총 {total_batches}개의 배치로 나누어 처리합니다.")

# 배치 단위로 처리
for batch_idx in tqdm(range(total_batches), desc="데이터 추가 진행"):
    # 현재 배치의 데이터 슬라이스
    start_idx = batch_idx * batch_size
    end_idx = min((batch_idx + 1) * batch_size, len(df))
    batch_df = df.iloc[start_idx:end_idx]
    
    # 1. 문서 준비
    documents = [
        f"제목: {row['title']}, 작가: {row['author']}, 소개글: {row['summary']}, 장르: {row['genre']}, 태그: {row['tags']}"
        for _, row in batch_df.iterrows()
    ]
    
    # 2. 메타데이터 준비
    metadatas = []
    for _, row in batch_df.iterrows():
        metadata = {
            "url": row["url"],
            "thumbnail": row["thumbnail"],
            "title": row["title"],
            "author": row["author"],
            "recommend": int(row["recommend"]),
            "genre": row["genre"],
            "end": bool(row["end"]),
            "publisher": row["publisher"] if pd.notna(row["publisher"]) else "",
            "page_count": int(row["page_count"]),
            "page_unit": row["page_unit"],
            "adult": bool(row["adult"]),
            "platform": row["platform"],
            "tags": row["tags"],
            "viewers": int(row["viewers"]),
            "type": row["type"],
            "rating": float(row["rating"]),
            "new_genre": row["new_genre"]
        }
        metadatas.append(metadata)
    
    # 3. ID 준비
    ids = [str(row["id"]) for _, row in batch_df.iterrows()]
    
    # 4. 현재 배치의 데이터 추가
    collection.add(
        documents=documents,
        metadatas=metadatas,
        ids=ids
    )
    
print(f'\n-- 3. 데이터 추가 완료: 총 {len(df)}개 작품 저장됨 --')
print(f"현재 컬렉션에 저장된 총 항목 수: {collection.count()}")

# 예: "로맨스 게임" 관련 소설 검색
print('\n-- 4. 검색 테스트: "로맨스 게임" --')
results = collection.query(
    query_texts=["로맨스 게임"],
    n_results=5
)
print(f"검색 결과 수: {len(results['documents'][0])}")
print("검색된 작품 제목:")
for i, doc_id in enumerate(results['ids'][0]):
    title = next((meta["title"] for meta in results['metadatas'][0] if meta), "제목 없음")
    print(f"{i+1}. {title} (ID: {doc_id})")

# 예: 특정 조건 필터링 (100만 이상 조회된 로맨스 소설)
print('\n-- 5. 필터링 검색: 100만 이상 조회된 로맨스 소설 --')
filtered_results = collection.query(
    query_texts=["로맨스"],
    filter={"viewers": {"$gt": 1000000}, "genre": "로맨스"},
    n_results=5
)
print(f"필터링 검색 결과 수: {len(filtered_results['documents'][0])}")
print("검색된 작품 제목:")
for i, doc_id in enumerate(filtered_results['ids'][0]):
    idx = filtered_results['ids'][0].index(doc_id)
    title = filtered_results['metadatas'][0][idx]['title']
    viewers = filtered_results['metadatas'][0][idx]['viewers']
    print(f"{i+1}. {title} (ID: {doc_id}, 조회수: {viewers:,})")