crjoya/bge-m3-proptech-retrieval

BGE-M3 PropTech Retriever

Korean property domain-specific retrieval model based on BGE-M3 light version

Dimensionality Reduction for a Lightweight Version of BGE-M3 Embeddings

• 고정된 bge embedding에 MLP 레이어 적용 → query에 pair된 positive 문서와 negative 문서 간 유사도 계산 → 상위 문서 추천
• Frozen Base Model + Trainable MLP Parameters: Triplet Loss for Fine-Tuning only fc layers

Dataset: 프롭테크 도메인 query-document 쌍 데이터셋

from datasets import load_dataset
from sklearn.model_selection import train_test_split
import numpy as np

# Load dataset
def load_data():
    dataset = load_dataset("crjoya/korean-proptech-retrieval")
    data_list = list(dataset['train'])
    train_data, eval_data = train_test_split(
            data_list, 
            test_size=0.2, 
            random_state=42
        )
    return train_data, eval_data

Usage

For testing simple document

from datasets import load_dataset
from huggingface_hub import hf_hub_download
from datasets import load_dataset
from sklearn.model_selection import train_test_split
import numpy as np
import torch
import json


# 1. Load models and data
def load_model(repo_id="crjoya/bge-m3-proptech-retrieval", device="cpu"):
    """Load models with consistent dtype handling"""
    model_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")

    # Initialize models with consistent dtypes
    # bge_model (base) 로드
    bge_model = BGEM3FlagModel("BAAI/bge-m3", use_fp16=False)  # Set to FP32
    mlp_model = BGE_M3_FineTuner().to(device).eval()
    mlp_model.load_state_dict(torch.load(model_path))

    return bge_model, mlp_model


def inference(query, documents, bge_model, mlp_model, device="cpu"):
    """Run inference with dtype handling"""
    # Get query embedding
    query_emb = bge_model.encode(query)["dense_vecs"]
    query_tensor = torch.tensor(query_emb, dtype=torch.float32).to(device)
    query_vec = mlp_model(query_tensor)

    # Get document embeddings
    similarities = []
    for doc in documents:
        doc_emb = bge_model.encode(doc)["dense_vecs"]
        doc_tensor = torch.tensor(doc_emb, dtype=torch.float32).to(device)
        # pass reduced dimension vector through MLP layers
        doc_vec = mlp_model(doc_tensor)
        sim = F.cosine_similarity(query_vec, doc_vec, dim=0)
        similarities.append(sim.item())

    return similarities


# Infer Usage
bge_model, mlp_model = load_model()

query = "서울 아파트 매매 시세 알아보고 싶어요"
documents = [
    "서울 강남구 아파트 매매가는 최근 평당 1억 원을 돌파했습니다.",
    "서울 아파트 전세 가격이 안정세를 보이고 있습니다.",
]

similarities = inference(query, documents, bge_model, mlp_model)