Views
No views yet
1from datasets import load_dataset
2from sklearn.model_selection import train_test_split
3import numpy as np
4
5# Load dataset
6def load_data():
7 dataset = load_dataset("crjoya/korean-proptech-retrieval")
8 data_list = list(dataset['train'])
9 train_data, eval_data = train_test_split(
10 data_list,
11 test_size=0.2,
12 random_state=42
13 )
14 return train_data, eval_data1from datasets import load_dataset
2from huggingface_hub import hf_hub_download
3from datasets import load_dataset
4from sklearn.model_selection import train_test_split
5import numpy as np
6import torch
7import json
8
9
10# 1. Load models and data
11def load_model(repo_id="crjoya/bge-m3-proptech-retrieval", device="cpu"):
12 """Load models with consistent dtype handling"""
13 model_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin")
14
15 # Initialize models with consistent dtypes
16 # bge_model (base) 로드
17 bge_model = BGEM3FlagModel("BAAI/bge-m3", use_fp16=False) # Set to FP32
18 mlp_model = BGE_M3_FineTuner().to(device).eval()
19 mlp_model.load_state_dict(torch.load(model_path))
20
21 return bge_model, mlp_model
22
23
24def inference(query, documents, bge_model, mlp_model, device="cpu"):
25 """Run inference with dtype handling"""
26 # Get query embedding
27 query_emb = bge_model.encode(query)["dense_vecs"]
28 query_tensor = torch.tensor(query_emb, dtype=torch.float32).to(device)
29 query_vec = mlp_model(query_tensor)
30
31 # Get document embeddings
32 similarities = []
33 for doc in documents:
34 doc_emb = bge_model.encode(doc)["dense_vecs"]
35 doc_tensor = torch.tensor(doc_emb, dtype=torch.float32).to(device)
36 # pass reduced dimension vector through MLP layers
37 doc_vec = mlp_model(doc_tensor)
38 sim = F.cosine_similarity(query_vec, doc_vec, dim=0)
39 similarities.append(sim.item())
40
41 return similarities
42
43
44# Infer Usage
45bge_model, mlp_model = load_model()
46
47query = "서울 아파트 매매 시세 알아보고 싶어요"
48documents = [
49 "서울 강남구 아파트 매매가는 최근 평당 1억 원을 돌파했습니다.",
50 "서울 아파트 전세 가격이 안정세를 보이고 있습니다.",
51]
52
53similarities = inference(query, documents, bge_model, mlp_model)