NMIXX-bge-icl

This repository contains a Bge-icl‐based Embedding model fine‐tuned with a triplet‐loss setup on the nmixx-fin/NMIXX_train dataset. It produces high‐quality sentence embeddings for Korean financial text, optimized for semantic similarity tasks in the finance domain.

---

How to use

import torch
import torch.nn.functional as F
from torch import Tensor
from transformers import AutoTokenizer, AutoModel

def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
    seq_lens = attention_mask.sum(dim=1) - 1
    idx = torch.arange(last_hidden_states.size(0), device=last_hidden_states.device)
    return last_hidden_states[idx, seq_lens]

def get_detailed_instruct(task: str, query: str) -> str:
    return f"<instruct>{task}\n<query>{query}"

def get_detailed_example(task: str, query: str, response: str) -> str:
    return f"<instruct>{task}\n<query>{query}\n<response>{response}"

def get_new_queries(queries, query_max_len, examples_prefix, tokenizer):
    tmp = tokenizer(
        queries,
        max_length=query_max_len - len(tokenizer("<s>", add_special_tokens=False)["input_ids"]) - len(tokenizer("\n<response></s>", add_special_tokens=False)["input_ids"]),
        truncation=True,
        return_tensors=None,
        add_special_tokens=False
    )
    prefix_ids = tokenizer(examples_prefix, add_special_tokens=False)["input_ids"]
    suffix_ids = tokenizer("\n<response>", add_special_tokens=False)["input_ids"]
    new_max = (len(prefix_ids) + len(suffix_ids) + query_max_len + 8) // 8 * 8 + 8
    decoded = tokenizer.batch_decode(tmp["input_ids"])
    return new_max, [examples_prefix + d + "\n<response>" for d in decoded]

model_name = "nmixx-fin/nmixx-bge-icl"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name).eval().to("cuda" if torch.cuda.is_available() else "cpu")

task = "제시된 기준 문장과 의미가 가장 유사한 문장을 찾으세요."
examples = [
    {
        "query": "나는 오늘 기분이 아주 좋아",
        "response": "오늘 정말 활기차고 행복한 하루였어요."
    },
    {
        "query": "바람이 많이 부는 날씨",
        "response": "바람이 세차게 불어 머리가 헝클어졌어요."
    }
]
example_strs = [get_detailed_example(task, e["query"], e["response"]) for e in examples]
examples_prefix = "\n\n".join(example_strs) + "\n\n"

queries = [
    get_detailed_instruct(task, "점심으로 피자를 먹었어요"),
    get_detailed_instruct(task, "비가 오려나?")
]
documents = [
    "오늘 햇빛이 쨍쨍해서 산책하기 딱 좋은 날씨였습니다.",
    "어제 저녁에 비가 내려서 길이 조금 젖어 있었습니다."
]

device = model.device
q_max, new_queries = get_new_queries(queries, 512, examples_prefix, tokenizer)

q_batch = tokenizer(new_queries, max_length=q_max, padding=True, truncation=True, return_tensors="pt").to(device)
d_batch = tokenizer(documents, max_length=512, padding=True, truncation=True, return_tensors="pt").to(device)

with torch.no_grad():
    q_out = model(**q_batch)
    q_emb = last_token_pool(q_out.last_hidden_state, q_batch["attention_mask"])
    d_out = model(**d_batch)
    d_emb = last_token_pool(d_out.last_hidden_state, d_batch["attention_mask"])

q_emb = F.normalize(q_emb, p=2, dim=1)
d_emb = F.normalize(d_emb, p=2, dim=1)

scores = (q_emb @ d_emb.T) * 100
print(scores.tolist())