from pathlib import Path

import torch
from datasets import Dataset
from sentence_transformers import CrossEncoder
from sentence_transformers.cross_encoder import (
    CrossEncoderTrainer,
    CrossEncoderTrainingArguments,
)
from sentence_transformers.cross_encoder.losses import BinaryCrossEntropyLoss


train_dataset = Dataset.from_dict(
    {
        "query": [
            "reset a user password",
            "reset a user password",
            "enable two-factor authentication",
            "enable two-factor authentication",
            "export audit logs",
            "export audit logs",
            "restore a deleted project",
            "restore a deleted project",
        ],
        "document": [
            "Open the user profile and send a password recovery email.",
            "Open billing settings and update the saved payment card.",
            "Open account security and enroll an authenticator app.",
            "Open the release dashboard and deploy the web service.",
            "Open compliance reports and export the audit log CSV file.",
            "Open team settings and invite a new member.",
            "Open deleted projects and restore the selected project.",
            "Open API tokens and create a replacement token.",
        ],
        "label": [1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0],
    }
)