searchKeybert.py

import os
import torch
import pandas as pd
import logging
import re
import faiss
import numpy as np
import time
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from datasets import load_dataset
from huggingface_hub import login
from sentence_transformers import SentenceTransformer
from joblib import Parallel, delayed
from tqdm import tqdm
from keybert import KeyBERT  # ✅ KeyBERT 추가

# ✅ 로그 설정
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

HF_API_TOKEN = os.getenv("HF_API_TOKEN")

if HF_API_TOKEN:
    try:
        logger.info("🔐 Hugging Face에 로그인 중...")
        login(token=HF_API_TOKEN)
        logger.info("✅ Hugging Face 로그인 성공!")
    except Exception as e:
        logger.error(f"❌ Hugging Face 로그인 실패: {e}")
        sys.exit("🚫 프로그램을 종료합니다. 유효한 HF_API_TOKEN이 필요합니다.")
else:
    logger.error("❌ 환경 변수 'HF_API_TOKEN'이 설정되지 않았습니다.")
    sys.exit("🚫 프로그램을 종료합니다. HF_API_TOKEN 환경 변수를 설정해 주세요.")

# ✅ FastAPI 인스턴스 생성
app = FastAPI(title="🚀 KeyBERT 기반 FAISS 검색 API", version="1.0")

# ✅ KeyBERT 모델 로드
embedding_model = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")
keyword_model = KeyBERT(embedding_model)  # ✅ KeyBERT 적용
logger.info("✅ KeyBERT 기반 키워드 추출 모델 로드 완료!")

# ✅ 진행 중인 경매 상품 데이터 로드
def load_huggingface_jsonl(dataset_name, split="train"):
    try:
        repo_id = f"aikobay/{dataset_name}"
        dataset = load_dataset(repo_id, split=split)
        df = dataset.to_pandas().dropna()
        return df
    except Exception as e:
        logger.error(f"❌ 데이터 로드 중 오류 발생: {e}")
        return pd.DataFrame()

try:
    active_sale_items = load_huggingface_jsonl("initial_saleitem_dataset")
    logger.info(f"✅ 진행 중인 경매 상품 데이터 로드 완료! 총 {len(active_sale_items)}개 상품")
except Exception as e:
    logger.error(f"❌ 상품 데이터 로드 중 오류 발생: {e}")
    active_sale_items = pd.DataFrame()

# ✅ FAISS 인덱스 초기화
faiss_index = faiss.IndexFlatL2(384)
indexed_items = []

# ✅ FAISS 인덱스 구축
def rebuild_faiss_index():
    global faiss_index, indexed_items, active_sale_items
    logger.info("🔄 새로운 sale_item 데이터로 FAISS 인덱스를 재구축합니다...")
    
    active_sale_items = load_huggingface_jsonl("initial_saleitem_dataset")
    item_names = active_sale_items["ITEMNAME"].tolist()
    indexed_items = item_names
    
    logger.info(f"🔹 총 {len(item_names)}개 상품 벡터화 시작...")
    item_vectors = embedding_model.encode(item_names, convert_to_numpy=True).astype("float32")
    
    faiss_index = faiss.IndexFlatL2(item_vectors.shape[1])
    faiss_index.add(item_vectors)
    logger.info(f"✅ FAISS 인덱스가 {len(indexed_items)}개 상품으로 새롭게 구축되었습니다.")

# ✅ KeyBERT 기반 키워드 추출 함수
def generate_similar_keywords(query: str, num_keywords: int = 5):
    """KeyBERT 모델을 이용해 검색어의 유사 키워드 추출"""
    try:
        keywords = keyword_model.extract_keywords(query, keyphrase_ngram_range=(1,2), stop_words=None, top_n=num_keywords)
        keywords = [kw[0] for kw in keywords]
        logger.info(f"🔍 생성된 유사 키워드: {keywords}")
        return keywords
    except Exception as e:
        logger.error(f"❌ KeyBERT 키워드 추출 중 오류 발생: {e}")
        return [query]

# ✅ FAISS 검색 함수 (유사 키워드 기반 검색)
def search_faiss_with_keywords(query: str, top_k: int = 5):
    start_time = time.time()
    keywords = generate_similar_keywords(query)
    keyword_vectors = embedding_model.encode(keywords, convert_to_numpy=True).astype("float32")
    
    all_results = []
    for vec in keyword_vectors:
        _, indices = faiss_index.search(np.array([vec]), top_k)
        all_results.extend(indices[0])
    
    unique_results = list(set(all_results))
    recommendations = [indexed_items[i] for i in unique_results]
    
    end_time = time.time()
    logger.info(f"🔍 검색 수행 완료! 걸린 시간: {end_time - start_time:.4f}초")
    return recommendations

# ✅ API 요청 모델
class RecommendRequest(BaseModel):
    search_query: str
    top_k: int = 5

# ✅ 추천 API 엔드포인트
@app.post("/api/recommend")
async def recommend(request: RecommendRequest):
    try:
        recommendations = search_faiss_with_keywords(request.search_query, request.top_k)
        return {"query": request.search_query, "recommendations": recommendations}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"추천 오류: {str(e)}")

# ✅ FAISS 인덱스 갱신 API
@app.post("/api/update_index")
async def update_index():
    rebuild_faiss_index()
    return {"message": "✅ FAISS 인덱스 업데이트 완료!"}

# ✅ FastAPI 실행
if __name__ == "__main__":
    rebuild_faiss_index()
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)