Arabic NLI Triplet - Sentence Transformer Model

This repository contains a fine-tuned Sentence Transformer model trained on the "Omartificial-Intelligence-Space/Arabic-NLi-Triplet" dataset. The model is trained to generate 384-dimensional embeddings for semantic similarity tasks like paraphrase mining, sentence similarity, and clustering in Arabic.

Model Overview

• Model Type: Sentence Transformer
• Base Model: intfloat/multilingual-e5-small
• Training Dataset: Omartificial-Intelligence-Space/Arabic-NLi-Triplet
• Similarity Function: Cosine Similarity
• Embedding Dimensionality: 384 dimensions
• Maximum Sequence Length: 128 tokens
• Performance Improvement: The model achieved around 10% improvement when tested on the test set of the provided dataset, compared to the base model's performance.

Dataset

Arabic NLI Triplet Dataset

The dataset contains triplets of sentences in Arabic: an anchor sentence, a positive sentence (semantically similar to the anchor), and a negative sentence (semantically dissimilar to the anchor). The dataset is designed for learning sentence representations through triplet margin loss.

Dataset Link: Omartificial-Intelligence-Space/Arabic-NLi-Triplet

Training Process

Loss Function: Triplet Margin Loss

We used the Triplet Margin Loss with a margin of 1.0. The model is trained to minimize the distance between anchor and positive embeddings, while maximizing the distance between anchor and negative embeddings.

Training Loss Progress:

Below is the training loss recorded at various steps during the training process:

Model Training Code

The model was trained using the following code (without resuming from checkpoints):

from datasets import load_dataset
from transformers import AutoTokenizer, AutoModel, TrainingArguments, Trainer
from torch.nn import TripletMarginLoss

# Load dataset
dataset = load_dataset("Omartificial-Intelligence-Space/Arabic-NLi-Triplet")

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("intfloat/multilingual-e5-small")

# Tokenize function
def tokenize_function(examples):
    anchor_encodings = tokenizer(examples['anchor'], truncation=True, padding='max_length', max_length=128)
    positive_encodings = tokenizer(examples['positive'], truncation=True, padding='max_length', max_length=128)
    negative_encodings = tokenizer(examples['negative'], truncation=True, padding='max_length', max_length=128)

    return {
        'anchor_input_ids': anchor_encodings['input_ids'],
        'anchor_attention_mask': anchor_encodings['attention_mask'],
        'positive_input_ids': positive_encodings['input_ids'],
        'positive_attention_mask': positive_encodings['attention_mask'],
        'negative_input_ids': negative_encodings['input_ids'],
        'negative_attention_mask': negative_encodings['attention_mask'],
    }

tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=dataset["train"].column_names)

# Define triplet loss
triplet_loss = TripletMarginLoss(margin=1.0)

def compute_loss(anchor_embedding, positive_embedding, negative_embedding):
    return triplet_loss(anchor_embedding, positive_embedding, negative_embedding)

# Load model
model = AutoModel.from_pretrained("intfloat/multilingual-e5-small")

class TripletTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        anchor_input_ids = inputs['anchor_input_ids'].to(self.args.device)
        anchor_attention_mask = inputs['anchor_attention_mask'].to(self.args.device)
        positive_input_ids = inputs['positive_input_ids'].to(self.args.device)
        positive_attention_mask = inputs['positive_attention_mask'].to(self.args.device)
        negative_input_ids = inputs['negative_input_ids'].to(self.args.device)
        negative_attention_mask = inputs['negative_attention_mask'].to(self.args.device)

        anchor_embeds = model(input_ids=anchor_input_ids, attention_mask=anchor_attention_mask).last_hidden_state.mean(dim=1)
        positive_embeds = model(input_ids=positive_input_ids, attention_mask=positive_attention_mask).last_hidden_state.mean(dim=1)
        negative_embeds = model(input_ids=negative_input_ids, attention_mask=negative_attention_mask).last_hidden_state.mean(dim=1)

        return compute_loss(anchor_embeds, positive_embeds, negative_embeds)

# Training arguments
training_args = TrainingArguments(
    output_dir="/content/drive/MyDrive/results",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir='/content/drive/MyDrive/logs',
    remove_unused_columns=False,
    fp16=True,
    save_total_limit=3,
)

# Initialize trainer
trainer = TripletTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets['train'],
)

# Start training
trainer.train()

# Save model and evaluate
trainer.save_model("/content/drive/MyDrive/fine-tuned-multilingual-e5")
results = trainer.evaluate()
print(results)

Framework Versions

• Python: 3.10.11
• Sentence Transformers: 3.0.1
• Transformers: 4.44.2
• PyTorch: 2.4.0
• Datasets: 2.21.0

How to Use

To use the model, install the required libraries and load the model with the following code:

pip install -U sentence-transformers

from sentence_transformers import SentenceTransformer

# Load the fine-tuned model
model = SentenceTransformer("gimmeursocks/ara-e5-small")

# Run inference
sentences = ['أنا سعيد', 'الجو جميل اليوم', 'هذا كلب كبير']
embeddings = model.encode(sentences)
print(embeddings.shape)

Citation

If you use this model or dataset, please cite the corresponding paper or dataset source.