Fine-tuning was performed using
PyLate, with contrastive training on the
rag-comprehensive-triplets dataset. It maps sentences & paragraphs to sequences of 128-dimensional dense vectors and can be used for semantic textual similarity using the MaxSim operator.
1import torch
2from pylate import models
3
4# Load the ColBERT model
5model = models.ColBERT("fjmgAI/col1-210M-EuroBERT", trust_remote_code=True)
6
7# Move the model to GPU if available, otherwise use CPU
8device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
9model.to(device)
10
11# Example data for similarity comparison
12query = "¿Cuál es la capital de España?" # Query sentence
13positive_doc = "La capital de España es Madrid." # Relevant document
14negative_doc = "Florida es un estado en los Estados Unidos." # Irrelevant document
15sentences = [query, positive_doc, negative_doc] # Combine all texts
16
17# Tokenize the input sentences using ColBERT's tokenizer
18inputs = model.tokenize(sentences)
19
20# Move all input tensors to the same device as the model (GPU/CPU)
21inputs = {key: value.to(device) for key, value in inputs.items()}
22
23# Generate token embeddings (no gradients needed for inference)
24with torch.no_grad():
25 embeddings_dict = model(inputs)
26 embeddings = embeddings_dict['token_embeddings']
27
28# Define ColBERT's MaxSim similarity function
29def colbert_similarity(query_emb, doc_emb):
30 """
31 Computes ColBERT-style similarity between query and document embeddings.
32 Uses maximum similarity (MaxSim) between individual tokens.
33
34 Args:
35 query_emb: [query_tokens, embedding_dim]
36 doc_emb: [doc_tokens, embedding_dim]
37
38 Returns:
39 Normalized similarity score
40 """
41 # Compute dot product between all token pairs
42 similarity_matrix = torch.matmul(query_emb, doc_emb.T)
43
44 # Get maximum similarity for each query token (MaxSim)
45 max_similarities = similarity_matrix.max(dim=1)[0]
46
47 # Return average of maximum similarities (normalized by query length)
48 return max_similarities.sum() / query_emb.shape[0]
49
50# Extract embeddings for each text
51query_emb = embeddings[0]
52positive_emb = embeddings[1]
53negative_emb = embeddings[2]
54
55# Compute similarity scores
56positive_score = colbert_similarity(query_emb, positive_emb)
57negative_score = colbert_similarity(query_emb, negative_emb)
58
59print(f"Similarity with positive document: {positive_score.item():.4f}")
60print(f"Similarity with negative document: {negative_score.item():.4f}")