Finetuning of BAAI/bge-m3 with 3GPP specification data.
Install the libraries
Install the FlagEmbedding via:
git clone https://github.com/FlagOpen/FlagEmbedding.git
cd FlagEmbedding
pip install -e .
or install without editing:
pip install -U FlagEmbedding
Inference
Dense Embedding
from FlagEmbedding import BGEM3FlagModel
model = BGEM3FlagModel('SadeghK/bge-m3-3gpp', use_fp16=True) # slight quality trade-off for speed
# queries (how an engineer might ask)
queries = [
"How is PDCP data recovery triggered for an AM DRB?",
"What is the purpose of the RRCReestablishment procedure?",
]
# candidate spec passages
passages = [
"For AM DRBs configured with PDCP data recovery, the PDCP entity performs retransmission of all PDCP SDUs previously submitted to lower layers, starting from the first SDU not acknowledged by RLC.",
"The RRC Reestablishment procedure re-establishes the RRC connection after a radio link failure, handover failure, or integrity check failure, restoring SRB1 and reactivating security.",
]
# match the training length (512), not 8192 — faster and in-distribution
q_emb = model.encode(queries, batch_size=12, max_length=512)['dense_vecs']
p_emb = model.encode(passages, batch_size=12, max_length=512)['dense_vecs']
similarity = q_emb @ p_emb.T
print(similarity) # diagonal (correct pairs) should dominate each row
# [[0.729 0.478 ], [0.452 0.7134]] -- The original BAAI/bge-m3
# [[0.579 0.0889], [0.1157 0.62 ]] -- The finetuned SadeghK/bge-m3-3gpp
Sparse Embedding (Lexical Weight)
from FlagEmbedding import BGEM3FlagModel
model = BGEM3FlagModel('SadeghK/bge-m3-3gpp', use_fp16=True) # slight quality trade-off for speed
queries = [
"How is PDCP data recovery triggered for an AM DRB?",
"What is the purpose of the RRCReestablishment procedure?",
]
passages = [
"For AM DRBs configured with PDCP data recovery, the PDCP entity retransmits all PDCP SDUs previously submitted to lower layers, starting from the first SDU not acknowledged by RLC.",
"The RRC Reestablishment procedure re-establishes the RRC connection after radio link failure, handover failure, or integrity check failure, restoring SRB1 and reactivating security.",
]
out_q = model.encode(queries, max_length=512, return_dense=True, return_sparse=True, return_colbert_vecs=False)
out_p = model.encode(passages, max_length=512, return_dense=True, return_sparse=True, return_colbert_vecs=False)
# inspect which tokens the sparse head weights (note the high mass on 3GPP identifiers)
print(model.convert_id_to_token(out_q['lexical_weights']))
# e.g. {'PD':0.20, 'CP':0.26, 'data':0.11, 'recovery':0.23, 'AM':0.14, 'DR':0.18, 'B':0.09, ...} (illustrative)
# full 2x2 lexical-matching matrix (rows = queries, cols = passages)
for i in range(len(queries)):
row = [model.compute_lexical_matching_score(out_q['lexical_weights'][i],
out_p['lexical_weights'][j])
for j in range(len(passages))]
print(row)
# diagonal (correct pairs) should be clearly larger than off-diagonal
# finetuned model SadeghK/bge-m3-3gpp
# [{'How': 0.0735, 'PD': 0.2104, 'CP': 0.275, 'data': 0.264, 'recovery': 0.3538, 'trigger': 0.323, 'ed': 0.005375, 'an': 0.01321, 'AM': 0.4426, 'DR': 0.211, 'B': 0.1913, '?': 0.04645}, {'What': 0.05386, 'is': 0.0368, 'the': 0.11115, 'purpose': 0.284, 'of': 0.0806, 'R': 0.01096, 'RC': 0.3489, 'Re': 0.2715, 'e': 0.326, 'stab': 0.3704, 'lish': 0.3208, 'ment': 0.2439, 'procedure': 0.38, '?': 0.068}]
# [0.667510986328125, 0], [0.002357959747314453, 0.7953815460205078]
# original model BAAI/bge-m3
# [{'How': 0.0578, 'is': 0.01768, 'PD': 0.149, 'CP': 0.2703, 'data': 0.1835, 'recovery': 0.2264, 'trigger': 0.221, 'ed': 0.03482, 'an': 0.00264, 'AM': 0.2323, 'DR': 0.1436, 'B': 0.1951, '?': 0.02673}, {'What': 0.04892, 'is': 0.03653, 'the': 0.02892, 'purpose': 0.1696, 'of': 0.01078, 'R': 0.10175, 'RC': 0.2703, 'Re': 0.1183, 'e': 0.07477, 'stab': 0.1932, 'lish': 0.1655, 'ment': 0.09375, 'procedure': 0.2163, '?': 0.0389}]
# [0.280487060546875, 0], [0.004868030548095703, 0.22907257080078125]
Multi-Vector (ColBERT)
from FlagEmbedding import BGEM3FlagModel
model = BGEM3FlagModel('SadeghK/bge-m3-3gpp', use_fp16=True)
queries = [
"How is PDCP data recovery triggered for an AM DRB?",
"What is the purpose of the RRCReestablishment procedure?",
]
passages = [
"For AM DRBs configured with PDCP data recovery, the PDCP entity retransmits all PDCP SDUs previously submitted to lower layers, starting from the first SDU not acknowledged by RLC.",
"The RRC Reestablishment procedure re-establishes the RRC connection after radio link failure, handover failure, or integrity check failure, restoring SRB1 and reactivating security.",
]
out_q = model.encode(queries, max_length=512,
return_dense=True, return_sparse=True, return_colbert_vecs=True)
out_p = model.encode(passages, max_length=512,
return_dense=True, return_sparse=True, return_colbert_vecs=True)
# full ColBERT matrix: rows = queries, cols = passages (diagonal = correct pairs)
for i in range(len(queries)):
row = [round(float(model.colbert_score(out_q['colbert_vecs'][i],
out_p['colbert_vecs'][j])), 4)
for j in range(len(passages))]
print(row)
# expect the diagonal (correct pairs) clearly above the off-diagonal
# [0.7694, 0.5045], [0.4794, 0.78] -- Original BAAI/bge-m3
# [0.7057, 0.2468], [0.2184, 0.7341] -- finetuned SadeghK/bge-m3-3gpp
Compute score for text pairs
Input a list of text pairs, you can get the scores computed by different methods.
from FlagEmbedding import BGEM3FlagModel
model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)
queries = [
"How is PDCP data recovery triggered for an AM DRB?",
"What is the purpose of the RRCReestablishment procedure?",
]
passages = [
"For AM DRBs configured with PDCP data recovery, the PDCP entity retransmits all PDCP SDUs previously submitted to lower layers, starting from the first SDU not acknowledged by RLC.",
"The RRC Reestablishment procedure re-establishes the RRC connection after radio link failure, handover failure, or integrity check failure, restoring SRB1 and reactivating security.",
]
sentence_pairs = [[q, p] for q in queries for p in passages]
scores = model.compute_score(
sentence_pairs,
max_query_length=512, # match training; queries are short so this is safe
max_passage_length=512, # lower (e.g. 128) for latency if your chunks are short
weights_for_different_modes=[0.4, 0.2, 0.4], # w[0]*dense + w[1]*sparse + w[2]*colbert
)
# readable view: label each pair (indices 0 & 3 are the correct matches)
labels = [f"q{i}->p{j}" for i in range(len(queries)) for j in range(len(passages))]
for mode in ["dense", "sparse", "colbert", "sparse+dense", "colbert+sparse+dense"]:
print(mode, {labels[k]: round(scores[mode][k], 4) for k in range(len(labels))})