pylate 다운로드 pylate 소스 코드 다운로드

파일 레이트

늦은 상호 작용 모델에 대한 유연한 훈련 및 검색

Pylate는 최첨단 Colbert 모델로 미세 조정, 추론 및 검색을 단순화하고 최적화하도록 설계된 문장 변압기 위에 제작 된 라이브러리입니다. 단일 및 다중 GPU 모두에서 쉽게 미세 조정할 수 있으므로 다양한 하드웨어 설정에 유연성을 제공합니다. Pylate는 또한 문서 검색을 간소화하고 광범위한 모델을로드 할 수 있으므로 대부분의 미리 훈련 된 언어 모델에서 Colbert 모델을 구성 할 수 있습니다.

설치

PIP를 사용하여 Pylate를 설치할 수 있습니다.

pip install pylate

평가 종속성의 경우 사용하십시오.

pip install " pylate[eval] "

선적 서류 비치

완전한 문서는 심층적 인 가이드, 예제 및 API 참조가 포함 된 여기에서 제공됩니다.

훈련

대조적 인 훈련

다음은 Pylate를 사용하여 MS Marco DataSet Triplet 데이터 세트에서 Colbert 모델을 교육하는 간단한 예입니다. 이 스크립트는 대조적 인 손실로 훈련을 보여주고 Hold-Out Eval 세트에서 모델을 평가합니다.

 import torch
from datasets import load_dataset
from sentence_transformers import (
    SentenceTransformerTrainer ,
    SentenceTransformerTrainingArguments ,
)

from pylate import evaluation , losses , models , utils

# Define model parameters for contrastive training
model_name = "bert-base-uncased"  # Choose the pre-trained model you want to use as base
batch_size = 32  # Larger batch size often improves results, but requires more memory

num_train_epochs = 1  # Adjust based on your requirements
# Set the run name for logging and output directory
run_name = "contrastive-bert-base-uncased"
output_dir = f"output/ { run_name } "

# 1. Here we define our ColBERT model. If not a ColBERT model, will add a linear layer to the base encoder.
model = models . ColBERT ( model_name_or_path = model_name )

# Compiling the model makes the training faster
model = torch . compile ( model )

# Load dataset
dataset = load_dataset ( "sentence-transformers/msmarco-bm25" , "triplet" , split = "train" )
# Split the dataset (this dataset does not have a validation set, so we split the training set)
splits = dataset . train_test_split ( test_size = 0.01 )
train_dataset = splits [ "train" ]
eval_dataset = splits [ "test" ]

# Define the loss function
train_loss = losses . Contrastive ( model = model )

# Initialize the evaluator
dev_evaluator = evaluation . ColBERTTripletEvaluator (
    anchors = eval_dataset [ "query" ],
    positives = eval_dataset [ "positive" ],
    negatives = eval_dataset [ "negative" ],
)

# Configure the training arguments (e.g., batch size, evaluation strategy, logging steps)
args = SentenceTransformerTrainingArguments (
    output_dir = output_dir ,
    num_train_epochs = num_train_epochs ,
    per_device_train_batch_size = batch_size ,
    per_device_eval_batch_size = batch_size ,
    fp16 = True ,  # Set to False if you get an error that your GPU can't run on FP16
    bf16 = False ,  # Set to True if you have a GPU that supports BF16
    run_name = run_name ,  # Will be used in W&B if `wandb` is installed
    learning_rate = 3e-6 ,
)

# Initialize the trainer for the contrastive training
trainer = SentenceTransformerTrainer (
    model = model ,
    args = args ,
    train_dataset = train_dataset ,
    eval_dataset = eval_dataset ,
    loss = train_loss ,
    evaluator = dev_evaluator ,
    data_collator = utils . ColBERTCollator ( model . tokenize ),
)
# Start the training process
trainer . train ()

훈련 후 출력 디렉토리 경로를 사용하여 모델을로드 할 수 있습니다.

 from pylate import models

model = models . ColBERT ( model_name_or_path = "contrastive-bert-base-uncased" )

지식 증류

Colbert 모델을 훈련 할 때 최상의 성능을 얻으려면 지식 증류를 사용하여 강력한 교사 모델의 점수를 사용하여 모델을 훈련시켜야합니다. 다음은 MS Marco의 Pylate에서 지식 증류를 사용하여 모델을 훈련시키는 방법에 대한 간단한 예입니다.

 import torch
from datasets import load_dataset
from sentence_transformers import (
    SentenceTransformerTrainer ,
    SentenceTransformerTrainingArguments ,
)

from pylate import losses , models , utils

# Load the datasets required for knowledge distillation (train, queries, documents)
train = load_dataset (
    path = "lightonai/ms-marco-en-bge" ,
    name = "train" ,
)

queries = load_dataset (
    path = "lightonai/ms-marco-en-bge" ,
    name = "queries" ,
)

documents = load_dataset (
    path = "lightonai/ms-marco-en-bge" ,
    name = "documents" ,
)

# Set the transformation to load the documents/queries texts using the corresponding ids on the fly
train . set_transform (
    utils . KDProcessing ( queries = queries , documents = documents ). transform ,
)

# Define the base model, training parameters, and output directory
model_name = "bert-base-uncased"  # Choose the pre-trained model you want to use as base
batch_size = 16
num_train_epochs = 1
# Set the run name for logging and output directory
run_name = "knowledge-distillation-bert-base"
output_dir = f"output/ { run_name } "

# Initialize the ColBERT model from the base model
model = models . ColBERT ( model_name_or_path = model_name )

# Compiling the model to make the training faster
model = torch . compile ( model )

# Configure the training arguments (e.g., epochs, batch size, learning rate)
args = SentenceTransformerTrainingArguments (
    output_dir = output_dir ,
    num_train_epochs = num_train_epochs ,
    per_device_train_batch_size = batch_size ,
    fp16 = True ,  # Set to False if you get an error that your GPU can't run on FP16
    bf16 = False ,  # Set to True if you have a GPU that supports BF16
    run_name = run_name ,
    learning_rate = 1e-5 ,
)

# Use the Distillation loss function for training
train_loss = losses . Distillation ( model = model )

# Initialize the trainer
trainer = SentenceTransformerTrainer (
    model = model ,
    args = args ,
    train_dataset = train ,
    loss = train_loss ,
    data_collator = utils . ColBERTCollator ( tokenize_fn = model . tokenize ),
)

# Start the training process
trainer . train ()

데이터 세트

Pylate는 포옹 페이스 데이터 세트를 지원하여 원활한 삼중 항 / 지식 증류 기반 교육을 가능하게합니다. 대조 교육을 위해 기존 문장 변압기 트리플렛 데이터 세트를 사용할 수 있습니다. 아래는 교육을위한 사용자 정의 삼중 항 데이터 세트를 작성하는 예입니다.

 from datasets import Dataset

dataset = [
    {
        "query" : "example query 1" ,
        "positive" : "example positive document 1" ,
        "negative" : "example negative document 1" ,
    },
    {
        "query" : "example query 2" ,
        "positive" : "example positive document 2" ,
        "negative" : "example negative document 2" ,
    },
    {
        "query" : "example query 3" ,
        "positive" : "example positive document 3" ,
        "negative" : "example negative document 3" ,
    },
]

dataset = Dataset . from_list ( mapping = dataset )

train_dataset , test_dataset = dataset . train_test_split ( test_size = 0.3 )

지식 증류 데이터 세트를 만들려면 다음 스 니펫을 사용할 수 있습니다.

 from datasets import Dataset

dataset = [
    {
        "query_id" : 54528 ,
        "document_ids" : [
            6862419 ,
            335116 ,
            339186 ,
        ],
        "scores" : [
            0.4546215673141326 ,
            0.6575686537173476 ,
            0.26825184192900203 ,
        ],
    },
    {
        "query_id" : 749480 ,
        "document_ids" : [
            6862419 ,
            335116 ,
            339186 ,
        ],
        "scores" : [
            0.2546215673141326 ,
            0.7575686537173476 ,
            0.96825184192900203 ,
        ],
    },
]


dataset = Dataset . from_list ( mapping = dataset )

documents = [
    { "document_id" : 6862419 , "text" : "example doc 1" },
    { "document_id" : 335116 , "text" : "example doc 2" },
    { "document_id" : 339186 , "text" : "example doc 3" },
]

queries = [
    { "query_id" : 749480 , "text" : "example query" },
]

documents = Dataset . from_list ( mapping = documents )

queries = Dataset . from_list ( mapping = queries )

검색하다

Pylate는 훈련 된 Colbert 모델 및 Voyager 색인을 사용하여 주어진 쿼리 세트에 대한 최상위 문서를 쉽게 검색 할 수 있습니다. 단순히 모델을로드하고 인덱스를 시작합니다.

 from pylate import indexes , models , retrieve

model = models . ColBERT (
    model_name_or_path = "lightonai/colbertv2.0" ,
)

index = indexes . Voyager (
    index_folder = "pylate-index" ,
    index_name = "index" ,
    override = True ,
)

retriever = retrieve . ColBERT ( index = index )

모델과 색인이 설정되면 임베딩 및 해당 ID를 사용하여 인덱스에 문서를 추가 할 수 있습니다.

 documents_ids = [ "1" , "2" , "3" ]

documents = [
    "document 1 text" , "document 2 text" , "document 3 text"
]

# Encode the documents
documents_embeddings = model . encode (
    documents ,
    batch_size = 32 ,
    is_query = False , # Encoding documents
    show_progress_bar = True ,
)

# Add the documents ids and embeddings to the Voyager index
index . add_documents (
    documents_ids = documents_ids ,
    documents_embeddings = documents_embeddings ,
)

그런 다음 주어진 쿼리 세트에 대한 Top-K 문서를 검색 할 수 있습니다.

 queries_embeddings = model . encode (
    [ "query for document 3" , "query for document 1" ],
    batch_size = 32 ,
    is_query = True , # Encoding queries
    show_progress_bar = True ,
)

scores = retriever . retrieve (
    queries_embeddings = queries_embeddings , 
    k = 10 ,
)

print ( scores )

샘플 출력 :

[
    [
        { "id" : "3" , "score" : 11.266985893249512 },
        { "id" : "1" , "score" : 10.303335189819336 },
        { "id" : "2" , "score" : 9.502392768859863 },
    ],
    [
        { "id" : "1" , "score" : 10.88800048828125 },
        { "id" : "3" , "score" : 9.950843811035156 },
        { "id" : "2" , "score" : 9.602447509765625 },
    ],
]

재진

Colbert 모델을 사용하여 인덱스를 구축하지 않고 1 단계 검색 파이프 라인 위에서 재평가를 수행하려면 Rank Function을 사용하여 쿼리 및 문서를 다시 전달할 수 있습니다.

 from pylate import rank

queries = [
    "query A" ,
    "query B" ,
]
documents = [
    [ "document A" , "document B" ],
    [ "document 1" , "document C" , "document B" ],
]
documents_ids = [
    [ 1 , 2 ],
    [ 1 , 3 , 2 ],
]

queries_embeddings = model . encode (
    queries ,
    is_query = True ,
)
documents_embeddings = model . encode (
    documents ,
    is_query = False ,
)

reranked_documents = rank . rerank (
    documents_ids = documents_ids ,
    queries_embeddings = queries_embeddings ,
    documents_embeddings = documents_embeddings ,
)

기여

우리는 기여를 환영합니다! 시작하려면 :

개발 종속성 설치 :

pip install " pylate[dev] "

실행 테스트 :

make test

Ruff와 함께 형식 코드 :

make ruff

문서 작성 :

make livedoc

소환

이 bibtex로 라이브러리를 참조 할 수 있습니다.

 @misc { PyLate ,
  title = { PyLate: Flexible Training and Retrieval for Late Interaction Models } ,
  author = { Chaffin, Antoine and Sourty, Raphaël } ,
  url = { https://github.com/lightonai/pylate } ,
  year = { 2024 }
}