from transformers import AutoTokenizer, AutoModel

# Tải bộ mã hóa để phân loại/nhúng
tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-encoder-150m")
model = AutoModel.from_pretrained("jhu-clsp/ettin-encoder-150m")

def predict_masked_token(text):
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    
    # Nhận dự đoán cho token [MASK]
    mask_indices = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)
    predictions = outputs.logits[mask_indices]
    
    # Nhận 5 dự đoán hàng đầu
    top_tokens = torch.topk(predictions, 5, dim=-1)
    return [tokenizer.decode(token) for token in top_tokens.indices[0]]

# Ví dụ
masked_text = "Thủ đô của Pháp là [MASK]."
predictions = predict_masked_token(masked_text)
print(f"Dự đoán: {predictions}")

from transformers import AutoTokenizer, AutoModelForCausalLM

# Tải bộ giải mã để tạo
tokenizer = AutoTokenizer.from_pretrained("jhu-clsp/ettin-decoder-150m")
model = AutoModelForCausalLM.from_pretrained("jhu-clsp/ettin-decoder-150m")

# Tạo văn bản
prompt = "Tương lai của trí tuệ nhân tạo là"
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(inputs.input_ids, max_length=50, temperature=0.7)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

import argparse

from datasets import load_dataset
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)
from sentence_transformers.evaluation import TripletEvaluator
from sentence_transformers.losses import CachedMultipleNegativesRankingLoss
from sentence_transformers.training_args import BatchSamplers

def main():
    # phân tích cú pháp lr & tên mô hình
    parser = argparse.ArgumentParser()
    parser.add_argument("--lr", type=float, default=8e-5)
    parser.add_argument("--model_name", type=str, default="jhu-clsp/ettin-encoder-150m")
    args = parser.parse_args()
    lr = args.lr
    model_name = args.model_name
    model_shortname = model_name.split("/")[-1]

    # 1. Tải mô hình để tinh chỉnh
    model = SentenceTransformer(model_name)

    # 2. Tải một tập dữ liệu để tinh chỉnh
    dataset = load_dataset(
        "sentence-transformers/msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1",
        "triplet-hard",
        split="train",
    )
    dataset_dict = dataset.train_test_split(test_size=1_000, seed=12)
    train_dataset = dataset_dict["train"].select(range(1_250_000))
    eval_dataset = dataset_dict["test"]

    # 3. Xác định hàm mất mát
    loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=16)  # Tăng mini_batch_size nếu bạn có đủ VRAM

    run_name = f"{model_shortname}-DPR-{lr}"
    # 4. (Tùy chọn) Chỉ định đối số đào tạo
    args = SentenceTransformerTrainingArguments(
        # Tham số bắt buộc:
        output_dir=f"output/{model_shortname}/{run_name}",
        # Tham số đào tạo tùy chọn:
        num_train_epochs=1,
        per_device_train_batch_size=512,
        per_device_eval_batch_size=512,
        warmup_ratio=0.05,
        fp16=False,  # Đặt thành False nếu GPU không thể xử lý FP16
        bf16=True,  # Đặt thành True nếu GPU hỗ trợ BF16
        batch_sampler=BatchSamplers.NO_DUPLICATES,  # (Cached) MultipleNegativesRankingLoss được hưởng lợi từ việc không có bản sao
        learning_rate=lr,
        # Các tham số theo dõi/gỡ lỗi tùy chọn:
        save_strategy="steps",
        save_steps=500,
        save_total_limit=2,
        logging_steps=500,
        run_name=run_name,  # Được sử dụng trong `wandb`, `tensorboard`, `neptune`, v.v. nếu được cài đặt
    )

    # 5. (Tùy chọn) Tạo bộ đánh giá & đánh giá mô hình cơ sở
    dev_evaluator = TripletEvaluator(
        anchors=eval_dataset["query"],
        positives=eval_dataset["positive"],
        negatives=eval_dataset["negative"],
        name="msmarco-co-condenser-dev",
    )
    dev_evaluator(model)

    # 6. Tạo một trình đào tạo & đào tạo
    trainer = SentenceTransformerTrainer(
        model=model,
        args=args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        loss=loss,
        evaluator=dev_evaluator,
    )
    trainer.train()

    # 7. (Tùy chọn) Đánh giá mô hình đã đào tạo trên bộ đánh giá sau khi đào tạo
    dev_evaluator(model)

    # 8. Lưu mô hình
    model.save_pretrained(f"output/{model_shortname}/{run_name}/final")

    # 9. (Tùy chọn) Đẩy nó lên Hugging Face Hub
    model.push_to_hub(run_name, private=False)

if __name__ == "__main__":
    main()

from datasets import load_dataset
from pylate import losses, models, utils
from sentence_transformers import (
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
)

def main():
    # Tải tập dữ liệu cần thiết cho việc chắt lọc kiến thức (đào tạo, truy vấn, tài liệu)
    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",
    )

    # Đặt biến đổi để tải văn bản tài liệu/truy vấn bằng cách sử dụng các id tương ứng một cách nhanh chóng
    train.set_transform(
        utils.KDProcessing(queries=queries, documents=documents).transform,
    )

    # Xác định mô hình cơ sở, tham số đào tạo và thư mục đầu ra
    num_train_epochs = 1
    lr = 8e-5
    batch_size = 16
    accum_steps = 1
    model_name = "jhu-clsp/ettin-encoder-150m"
    model_shortname = model_name.split("/")[-1]

    # Đặt tên chạy cho ghi nhật ký và thư mục đầu ra
    run_name = f"{model_shortname}-colbert-KD-{lr}"
    output_dir = f"output/{model_shortname}/{run_name}"

    # Khởi tạo mô hình ColBERT từ mô hình cơ sở
    model = models.ColBERT(model_name_or_path=model_name)

    # Định cấu hình các đối số đào tạo (ví dụ: số kỷ nguyên, kích thước lô, tốc độ học)
    args = SentenceTransformerTrainingArguments(
        output_dir=output_dir,
        num_train_epochs=num_train_epochs,
        per_device_train_batch_size=batch_size,
        fp16=False,  # Đặt thành False nếu bạn gặp lỗi rằng GPU của bạn không thể chạy trên FP16
        bf16=True,  # Đặt thành True nếu bạn có GPU hỗ trợ BF16
        run_name=run_name,
        logging_steps=10,
        learning_rate=lr,
        gradient_accumulation_steps=accum_steps,
        warmup_ratio=0.05,
    )

    # Sử dụng hàm mất mát Distillation để đào tạo
    train_loss = losses.Distillation(model=model)

    # Khởi tạo trình đào tạo
    trainer = SentenceTransformerTrainer(
        model=model,
        args=args,
        train_dataset=train,
        loss=train_loss,
        data_collator=utils.ColBERTCollator(tokenize_fn=model.tokenize),
    )

    # Bắt đầu quá trình đào tạo
    trainer.train()

    model.save_pretrained(f"{output_dir}/final")

if __name__ == "__main__":
    main()

import logging

from datasets import load_dataset

from sentence_transformers import (
    SparseEncoder,
    SparseEncoderModelCardData,
    SparseEncoderTrainer,
    SparseEncoderTrainingArguments,
)
from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
from sentence_transformers.sparse_encoder.losses import SparseMultipleNegativesRankingLoss, SpladeLoss
from sentence_transformers.training_args import BatchSamplers

logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)

# 1. Tải một mô hình để tinh chỉnh với 2. (Tùy chọn) dữ liệu thẻ mô hình
model = SparseEncoder(
    "jhu-clsp/ettin-encoder-150m",
    model_card_data=SparseEncoderModelCardData(
        language="en",
        license="apache-2.0",
    )
)

# 3. Tải một tập dữ liệu để tinh chỉnh
full_dataset = load_dataset("sentence-transformers/natural-questions", split="train").select(range(100_000))
dataset_dict = full_dataset.train_test_split(test_size=1_000, seed=12)
train_dataset = dataset_dict["train"]
eval_dataset = dataset_dict["test"]

# 4. Xác định hàm mất mát
loss = SpladeLoss(
    model=model,
    loss=SparseMultipleNegativesRankingLoss(model=model),
    query_regularizer_weight=5e-5,
    document_regularizer_weight=3e-5,
)

# 5. (Tùy chọn) Chỉ định đối số đào tạo
run_name = "splade-distilbert-base-uncased-nq"
args = SparseEncoderTrainingArguments(
    # Tham số bắt buộc:
    output_dir=f"models/{run_name}",
    # Tham số đào tạo tùy chọn:
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    learning_rate=2e-5,
    warmup_ratio=0.1,
    fp16=True,  # Đặt thành False nếu bạn gặp lỗi rằng GPU của bạn không thể chạy trên FP16
    bf16=False,  # Đặt thành True nếu bạn có GPU hỗ trợ BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # MultipleNegativesRankingLoss được hưởng lợi từ việc không có mẫu trùng lặp trong một lô
    # Các tham số theo dõi/gỡ lỗi tùy chọn:
    eval_strategy="steps",
    eval_steps=1000,
    save_strategy="steps",
    save_steps=1000,
    save_total_limit=2,
    logging_steps=200,
    run_name=run_name,  # Sẽ được sử dụng trong W&B nếu `wandb` được cài đặt
)

# 6. (Tùy chọn) Tạo bộ đánh giá & đánh giá mô hình cơ sở
dev_evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"], batch_size=16)

# 7. Tạo một trình đào tạo & đào tạo
trainer = SparseEncoderTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss=loss,
    evaluator=dev_evaluator,
)
trainer.train()

# 8. Đánh giá lại hiệu suất của mô hình sau khi đào tạo
dev_evaluator(model)

# 9. Lưu mô hình đã đào tạo
model.save_pretrained(f"models/{run_name}/final")

# 10. (Tùy chọn) Đẩy nó lên Hugging Face Hub
model.push_to_hub(run_name)

import logging
import traceback

import torch
from datasets import load_dataset

from sentence_transformers import SentenceTransformer
from sentence_transformers.cross_encoder import (
    CrossEncoder,
    CrossEncoderModelCardData,
    CrossEncoderTrainer,
    CrossEncoderTrainingArguments,
)
from sentence_transformers.cross_encoder.evaluation import (
    CrossEncoderNanoBEIREvaluator,
    CrossEncoderRerankingEvaluator,
)
from sentence_transformers.cross_encoder.losses import BinaryCrossEntropyLoss
from sentence_transformers.evaluation import SequentialEvaluator
from sentence_transformers.util import mine_hard_negatives

# Đặt mức nhật ký thành INFO để nhận thêm thông tin
logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)


def main():
    model_name = "jhu-clsp/ettin-encoder-150m"

    train_batch_size = 64
    num_epochs = 1
    num_hard_negatives = 5  # Cần khai thác bao nhiêu số âm cứng cho mỗi cặp câu hỏi-trả lời

    # 1a. Tải một mô hình để tinh chỉnh với 1b. (Tùy chọn) dữ liệu thẻ mô hình
    model = CrossEncoder(
        model_name,
        model_card_data=CrossEncoderModelCardData(
            language="en",
            license="apache-2.0",
        ),
    )
    print("Độ dài tối đa của mô hình:", model.max_length)
    print("Số nhãn của mô hình:", model.num_labels)

    # 2a. Tải tập dữ liệu GooAQ: https://huggingface.co/datasets/sentence-transformers/gooaq
    logging.info("Đọc tập dữ liệu đào tạo gooaq")
    full_dataset = load_dataset("sentence-transformers/gooaq", split="train").select(range(100_000))
    dataset_dict = full_dataset.train_test_split(test_size=1_000, seed=12)
    train_dataset = dataset_dict["train"]
    eval_dataset = dataset_dict["test"]
    logging.info(train_dataset)
    logging.info(eval_dataset)

    # 2b. Sửa đổi tập dữ liệu đào tạo của chúng tôi để bao gồm các số âm cứng bằng cách sử dụng mô hình nhúng rất hiệu quả
    embedding_model = SentenceTransformer("sentence-transformers/static-retrieval-mrl-en-v1", device="cpu")
    hard_train_dataset = mine_hard_negatives(
        train_dataset,
        embedding_model,
        num_negatives=num_hard_negatives,  # Bao nhiêu số âm trên mỗi cặp câu hỏi-trả lời
        margin=0,  # Độ tương đồng giữa truy vấn và mẫu âm phải thấp hơn độ tương đồng truy vấn-dương
        range_min=0,  # Bỏ qua x mẫu tương tự nhất
        range_max=100,  # Chỉ xem xét x mẫu tương tự nhất
        sampling_strategy="top",  # Lấy mẫu các số âm hàng đầu từ phạm vi
        batch_size=4096,  # Sử dụng kích thước lô là 4096 cho mô hình nhúng
        output_format="labeled-pair",  # Định dạng đầu ra là (truy vấn, đoạn văn, nhãn), như BinaryCrossEntropyLoss yêu cầu
        use_faiss=True,
    )
    logging.info(hard_train_dataset)

    # 2c. (Tùy chọn) Lưu tập dữ liệu đào tạo cứng vào đĩa
    # hard_train_dataset.save_to_disk("gooaq-hard-train")
    # Tải lại với:
    # hard_train_dataset = load_from_disk("gooaq-hard-train")

    # 3. Xác định mất mát đào tạo của chúng tôi.
    # Nên đặt pos_weight là tỷ lệ giữa số dương với số âm, a.k.a. `num_hard_negatives`
    loss = BinaryCrossEntropyLoss(model=model, pos_weight=torch.tensor(num_hard_negatives))

    # 4a. Xác định bộ đánh giá. Chúng tôi sử dụng CrossEncoderNanoBEIREvaluator, là một bộ đánh giá nhẹ cho việc xếp hạng lại tiếng Anh
    nano_beir_evaluator = CrossEncoderNanoBEIREvaluator(
        dataset_names=["msmarco", "nfcorpus", "nq"],
        batch_size=train_batch_size,
    )

    # 4b. Xác định một trình đánh giá xếp hạng lại bằng cách khai thác các số âm cứng cho các cặp câu hỏi-trả lời
    # Chúng tôi bao gồm câu trả lời tích cực trong danh sách các số âm, để bộ đánh giá có thể sử dụng hiệu suất của
    # mô hình nhúng làm đường cơ sở.
    hard_eval_dataset = mine_hard_negatives(
        eval_dataset,
        embedding_model,
        corpus=full_dataset["answer"],  # Sử dụng toàn bộ tập dữ liệu làm ngữ liệu
        num_negatives=30,  # Cần xếp hạng lại bao nhiêu tài liệu
        batch_size=4096,
        include_positives=True,
        output_format="n-tuple",
        use_faiss=True,
    )
    logging.info(hard_eval_dataset)
    reranking_evaluator = CrossEncoderRerankingEvaluator(
        samples=[
            {
                "query": sample["question"],
                "positive": [sample["answer"]],
                "documents": [sample[column_name] for column_name in hard_eval_dataset.column_names[2:]],
            }
            for sample in hard_eval_dataset
        ],
        batch_size=train_batch_size,
        name="gooaq-dev",
        # Cài đặt thực tế: chỉ xếp hạng lại các số dương mà trình truy xuất đã tìm thấy
        # Đặt thành True để xếp hạng lại *tất cả* các số dương
        always_rerank_positives=False,
    )

    # 4c. Kết hợp bộ đánh giá & chạy mô hình cơ sở trên chúng
    evaluator = SequentialEvaluator([reranking_evaluator, nano_beir_evaluator])
    evaluator(model)

    # 5. Xác định các đối số đào tạo
    short_model_name = model_name if "/" not in model_name else model_name.split("/")[-1]
    run_name = f"reranker-{short_model_name}-gooaq-bce"
    args = CrossEncoderTrainingArguments(
        # Tham số bắt buộc:
        output_dir=f"models/{run_name}",
        # Tham số đào tạo tùy chọn:
        num_train_epochs=num_epochs,
        per_device_train_batch_size=train_batch_size,
        per_device_eval_batch_size=train_batch_size,
        learning_rate=2e-5,
        warmup_ratio=0.1,
        fp16=False,  # Đặt thành False nếu bạn gặp lỗi rằng GPU của bạn không thể chạy trên FP16
        bf16=True,  # Đặt thành True nếu bạn có GPU hỗ trợ BF16
        dataloader_num_workers=4,
        load_best_model_at_end=True,
        metric_for_best_model="eval_gooaq-dev_ndcg@10",
        # Các tham số theo dõi/gỡ lỗi tùy chọn:
        eval_strategy="steps",
        eval_steps=1000,
        save_strategy="steps",
        save_steps=1000,
        save_total_limit=2,
        logging_steps=200,
        logging_first_step=True,
        run_name=run_name,  # Sẽ được sử dụng trong W&B nếu `wandb` được cài đặt
        seed=12,
    )

    # 6. Tạo trình đào tạo & bắt đầu đào tạo
    trainer = CrossEncoderTrainer(
        model=model,
        args=args,
        train_dataset=hard_train_dataset,
        loss=loss,
        evaluator=evaluator,
    )
    trainer.train()

    # 7. Đánh giá mô hình cuối cùng, hữu ích để bao gồm các mô hình này trong thẻ mô hình
    evaluator(model)

    # 8. Lưu mô hình cuối cùng
    final_output_dir = f"models/{run_name}/final"
    model.save_pretrained(final_output_dir)

    # 9. (Tùy chọn) lưu mô hình vào Hugging Face Hub!
    # Nên chạy `huggingface-cli login` để đăng nhập vào tài khoản Hugging Face của bạn trước
    try:
        model.push_to_hub(run_name)
    except Exception:
        logging.error(
            f"Lỗi tải lên mô hình lên Hugging Face Hub:\n{t

### [Link bài viết gốc](https://huggingface.co/blog/ettin)