yjoonjang/splade-ko-v1

SparseEncoder(
  (0): MLMTransformer({'max_seq_length': 8192, 'do_lower_case': False, 'architecture': 'ModernBertForMaskedLM'})
  (1): SpladePooling({'pooling_strategy': 'max', 'activation_function': 'relu', 'word_embedding_dimension': 50000})
)

pip install -U sentence-transformers

from sentence_transformers import SparseEncoder

# Download from the 🤗 Hub
model = SparseEncoder("yjoonjang/splade-ko-v1")
# Run inference
sentences = [
    '양이온 최적화 방법은 산소공공을 감소시키기 때문에 전자 농도가 증가하는 문제점을 갖고있을까?',
    '산화물 TFT 소자 신뢰성 열화기구\n그러나 이와 같은 양이온 최적화 방법은 산소공공을 감소시키기 때문에 전자농도 역시 감소하게 되어 전계 이동도가 감소하는 문제점을 않고 있다.\n이는 산화물 반도체의 전도기구가 Percolation Conduction에 따르기 때문이다. ',
    '세포대사 기능 분석을 위한 광학센서 기반 용존산소와 pH 측정 시스템의 제작 및 특성 분석\n수소이온 농도가 증가하는 경우인 \\( \\mathrm{pH} \\) 가 낮아지면 다수의 수소이온들과 충돌한 방출 광이 에너지를 잃고 짧은 검출시간을 갖는다. \n반대로 \\( \\mathrm{pH} \\)가 높아질수록 형광물질로부터 방출된 광의 수명이 길어져 긴 검출시간을 가진다.',
]
embeddings = model.encode(sentences)
print(embeddings.shape)
# [3, 50000]

# Get the similarity scores for the embeddings
similarities = model.similarity(embeddings, embeddings)
print(similarities)
# tensor([[ 46.0239,  57.8961,  22.8014],
#         [ 57.8961, 270.6235,  56.5666],
#         [ 22.8014,  56.5666, 275.8828]], device='cuda:0')

  {
  	"query": "Is it unopened?",
  	"document": "No. It is a renewed product."
  },
  {
  	"query": "Is it compatible with iPad Air 3?",
  	"document": "Yes, it is possible."
  }

from sentence_transformers import SparseEncoder
from datasets import load_dataset
from sentence_transformers.sparse_encoder.evaluation import SparseInformationRetrievalEvaluator
import os
import pandas as pd
from tqdm import tqdm
import json
from multiprocessing import Process, current_process
import torch
from setproctitle import setproctitle
import traceback

# GPU별로 평가할 데이터셋 매핑
DATASET_GPU_MAPPING = {
	0: [
		"yjoonjang/markers_bm",
		"taeminlee/Ko-StrategyQA",
		"facebook/belebele",
		"xhluca/publichealth-qa",
		"Shitao/MLDR"
	],
	1: [
		"miracl/mmteb-miracl",
	],
	2: [
		"mteb/mrtidy",
	]
}

model_name = "yjoonjang/splade-ko-v1"

def evaluate_dataset(model_name, gpu_id, eval_datasets, output_dir):
	output_dir = f"{output_dir}/{model_name}"
	os.makedirs(output_dir, exist_ok=True)
	
	"""단일 GPU에서 할당된 데이터셋들을 평가하는 함수"""
	import torch
	try:
		device = torch.device(f"cuda:{str(gpu_id)}") 
		torch.cuda.set_device(device)
		
		setproctitle(f"yjoonjang splade-eval-gpu{gpu_id}")
		print(f"Running datasets: {eval_datasets} on GPU {gpu_id} in process {current_process().name}")
		
		# 모델 로드
		model = SparseEncoder(model_name, trust_remote_code=True, device=device)
		
		for eval_dataset in eval_datasets:
			short_dataset_name = eval_dataset.split("/")[-1]
			
			prediction_filepath = f"{output_dir}/{short_dataset_name}.json"
			if os.path.exists(prediction_filepath):
				print(f"Skipping evaluation for {eval_dataset} as output already exists at {prediction_filepath}")
				continue

			corpus = {}
			queries = {}
			relevant_docs = {}
			split = "dev"
			if eval_dataset == "yjoonjang/markers_bm" or eval_dataset == "yjoonjang/squad_kor_v1":
				split = "test"

			if eval_dataset in ["yjoonjang/markers_bm", "taeminlee/Ko-StrategyQA"]:
				dev_corpus = load_dataset(eval_dataset, "corpus", split="corpus")
				dev_queries = load_dataset(eval_dataset, "queries", split="queries")
				relevant_docs_data = load_dataset(eval_dataset, "default", split=split)
				
				queries = dict(zip(dev_queries["_id"], dev_queries["text"]))
				# Combine title and text if title exists (MTEB format)
				if "title" in dev_corpus.column_names and "text" in dev_corpus.column_names:
					corpus = {
						row["_id"]: (row["title"] + " " + row["text"]).strip()
						for row in dev_corpus
					}
				elif "text" in dev_corpus.column_names:
					corpus = dict(zip(dev_corpus["_id"], dev_corpus["text"]))
				else:
					raise ValueError(f"Corpus dataset must have 'text' field")
				for row in relevant_docs_data:
					qid_str = str(row["query-id"])
					corpus_ids_str = str(row["corpus-id"])
					score = row.get("score", 1)  # Default to 1 if no score field
					if score >= 1:  # Only include relevant documents (score >= 1)
						if qid_str not in relevant_docs:
							relevant_docs[qid_str] = set()
						relevant_docs[qid_str].add(corpus_ids_str)

			elif eval_dataset == "facebook/belebele":
				split = "test"
				ds = load_dataset(eval_dataset, "kor_Hang", split=split)
				
				corpus_df = pd.DataFrame(ds)
				corpus_df = corpus_df.drop_duplicates(subset=["link"])
				corpus_df["cid"] = [f"C{i}" for i in range(len(corpus_df))]
				corpus = dict(zip(corpus_df["cid"], corpus_df["flores_passage"]))
				
				link_to_cid = dict(zip(corpus_df["link"], corpus_df["cid"]))
				
				queries_df = pd.DataFrame(ds)
				queries_df = queries_df.drop_duplicates(subset=["question"])
				queries_df["qid"] = [f"Q{i}" for i in range(len(queries_df))]
				queries = dict(zip(queries_df["qid"], queries_df["question"]))
				
				question_to_qid = dict(zip(queries_df["question"], queries_df["qid"]))

				for row in tqdm(ds, desc="Processing belebele"):
					qid = question_to_qid[row["question"]]
					cid = link_to_cid[row["link"]]
					if qid not in relevant_docs:
						relevant_docs[qid] = set()
					relevant_docs[qid].add(cid)

			elif eval_dataset == "miracl/mmteb-miracl":
				split = "dev"
				corpus_ds = load_dataset(eval_dataset, "corpus-ko", split="corpus", trust_remote_code=True)
				queries_ds = load_dataset(eval_dataset, "queries-ko", split="queries", trust_remote_code=True)
				qrels_ds = load_dataset(eval_dataset, "ko", split=split, trust_remote_code=True)

				# Combine title and text if title exists (MTEB format)
				if "title" in corpus_ds.column_names:
					corpus = {
						row['docid']: (row['title'] + " " + row['text']).strip()
						for row in corpus_ds
					}
				else:
					corpus = {row['docid']: row['text'] for row in corpus_ds}
				queries = {row['query_id']: row['query'] for row in queries_ds}

				for row in qrels_ds:
					qid = row["query_id"]
					cid = row["docid"]
					score = row.get("score", 1)  # Default to 1 if no score field
					if score >= 1:  # Only include relevant documents (score >= 1)
						if qid not in relevant_docs:
							relevant_docs[qid] = set()
						relevant_docs[qid].add(cid)

			elif eval_dataset == "mteb/mrtidy":
				split = "test"
				corpus_ds = load_dataset(eval_dataset, "korean-corpus", split="train", trust_remote_code=True)
				queries_ds = load_dataset(eval_dataset, "korean-queries", split=split, trust_remote_code=True)
				qrels_ds = load_dataset(eval_dataset, "korean-qrels", split=split, trust_remote_code=True)

				# Combine title and text if title exists (MTEB format)
				if "title" in corpus_ds.column_names and "text" in corpus_ds.column_names:
					corpus = {
						row["_id"]: (row["title"] + " " + row["text"]).strip()
						for row in corpus_ds
					}
				elif "text" in corpus_ds.column_names:
					corpus = {row['_id']: row['text'] for row in corpus_ds}
				else:
					raise ValueError(f"Corpus dataset must have 'text' field")
				queries = {row['_id']: row['text'] for row in queries_ds}

				for row in qrels_ds:
					qid = str(row["query-id"])
					cid = str(row["corpus-id"])
					score = row.get("score", 1)  # Default to 1 if no score field
					if score >= 1:  # Only include relevant documents (score >= 1)
						if qid not in relevant_docs:
							relevant_docs[qid] = set()
						relevant_docs[qid].add(cid)

			elif eval_dataset == "Shitao/MLDR":
				split = "dev"
				corpus_ds = load_dataset(eval_dataset, "corpus-ko", split="corpus", trust_remote_code=True)
				lang_data = load_dataset(eval_dataset, "ko", split=split, trust_remote_code=True)
				
				# Combine title and text if title exists (MTEB format)
				if "title" in corpus_ds.column_names:
					corpus = {
						row['docid']: (row['title'] + " " + row['text']).strip()
						for row in corpus_ds
					}
				else:
					corpus = {row['docid']: row['text'] for row in corpus_ds}
				queries = {row['query_id']: row['query'] for row in lang_data}

				for row in lang_data:
					qid = row["query_id"]
					cid = row["positive_passages"][0]["docid"]
					if qid not in relevant_docs:
						relevant_docs[qid] = set()
					relevant_docs[qid].add(cid)

			elif eval_dataset == "xhluca/publichealth-qa":
				split = "test"
				ds = load_dataset(eval_dataset, "korean", split=split, trust_remote_code=True)
				
				ds = ds.filter(lambda x: x["question"] is not None and x["answer"] is not None)
				
				corpus_df = pd.DataFrame(list(ds))
				corpus_df = corpus_df.drop_duplicates(subset=["answer"])
				corpus_df["cid"] = [f"D{i}" for i in range(len(corpus_df))]
				corpus = dict(zip(corpus_df["cid"], corpus_df["answer"]))
				answer_to_cid = dict(zip(corpus_df["answer"], corpus_df["cid"]))

				queries_df = pd.DataFrame(list(ds))
				queries_df = queries_df.drop_duplicates(subset=["question"])
				queries_df["qid"] = [f"Q{i}" for i in range(len(queries_df))]
				queries = dict(zip(queries_df["qid"], queries_df["question"]))
				question_to_qid = dict(zip(queries_df["question"], queries_df["qid"]))
				
				for row in tqdm(ds, desc="Processing publichealth-qa"):
					qid = question_to_qid[row["question"]]
					cid = answer_to_cid[row["answer"]]
					if qid not in relevant_docs:
						relevant_docs[qid] = set()
					relevant_docs[qid].add(cid)

			else:
				continue
			
			if torch.cuda.get_device_name().startswith('NVIDIA A100'):
				batch_size = 16
			else:
				batch_size = 4
					
			evaluator = SparseInformationRetrievalEvaluator(
				queries=queries,
				corpus=corpus,
				relevant_docs=relevant_docs,
				write_csv=False,
				name=f"{eval_dataset}",
				show_progress_bar=True,
				batch_size=batch_size,
				write_predictions=False
			)
			short_dataset_name = eval_dataset.split("/")[-1]
			output_filepath = f"{output_dir}/{short_dataset_name}.json"
			metrics = evaluator(model)
			print(f"GPU {gpu_id} - {eval_dataset} metrics: {metrics}")
			with open(output_filepath, "w", encoding="utf-8") as f:
				json.dump(metrics, f, ensure_ascii=False, indent=2)
				
	except Exception as ex:
		print(f"Error on GPU {gpu_id}: {ex}")
		traceback.print_exc()

if __name__ == "__main__":
	torch.multiprocessing.set_start_method('spawn')
	
	print(f"Starting evaluation for model: {model_name}")
	output_dir = "./results_inference_free_new_idf"
	# output_dir = "./results_inference_free_new"
	processes = []
	
	for gpu_id, datasets in DATASET_GPU_MAPPING.items():
		p = Process(target=evaluate_dataset, args=(model_name, gpu_id, datasets, output_dir))
		p.start()
		processes.append(p)
	
	for p in processes:
		p.join()
	
	print(f"Completed evaluation for model: {model_name}")

@software{spladeko,
  author       = {Youngjoon, Jang},
  title        = {splade-ko-v1},
  subtitle     = {Independent Korean SPLADE Model for the Industry},
  publisher    = {Hugging Face},
  year         = {2025},
  version      = {1.0.0},
  url          = {https://huggingface.co/yjoonjang/splade-ko-v1}
}

@inproceedings{reimers-2019-sentence-bert,
    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
    author = "Reimers, Nils and Gurevych, Iryna",
    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
    month = "11",
    year = "2019",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/1908.10084",
}

@misc{formal2022distillationhardnegativesampling,
      title={From Distillation to Hard Negative Sampling: Making Sparse Neural IR Models More Effective},
      author={Thibault Formal and Carlos Lassance and Benjamin Piwowarski and Stéphane Clinchant},
      year={2022},
      eprint={2205.04733},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2205.04733},
}

@misc{henderson2017efficient,
    title={Efficient Natural Language Response Suggestion for Smart Reply},
    author={Matthew Henderson and Rami Al-Rfou and Brian Strope and Yun-hsuan Sung and Laszlo Lukacs and Ruiqi Guo and Sanjiv Kumar and Balint Miklos and Ray Kurzweil},
    year={2017},
    eprint={1705.00652},
    archivePrefix={arXiv},
    primaryClass={cs.CL}
}

@article{paria2020minimizing,
    title={Minimizing flops to learn efficient sparse representations},
    author={Paria, Biswajit and Yeh, Chih-Kuan and Yen, Ian EH and Xu, Ning and Ravikumar, Pradeep and P{'o}czos, Barnab{'a}s},
    journal={arXiv preprint arXiv:2004.05665},
    year={2020}
}