"""Train the intent classifier offline (for example in Colab), never at API startup."""
from pathlib import Path

LABELS = ["MATCH_STATISTICS", "MATCH_RECAP", "MOMENTUM_ANALYSIS", "PLAYER_COMPARISON", "CRICKET_RULES", "GENERAL_MATCH_QUERY"]


def train(output_dir: Path, dataset_path: Path) -> None:
    from datasets import load_dataset
    from peft import LoraConfig, TaskType, get_peft_model
    from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments

    model_name = "google/bert_uncased_L-4_H-256_A-4"
    dataset = load_dataset("json", data_files=str(dataset_path))["train"].train_test_split(test_size=0.2, seed=42)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    label_ids = {label: index for index, label in enumerate(LABELS)}
    tokenized = dataset.map(
        lambda row: {**tokenizer(row["text"], truncation=True), "labels": [label_ids[value] for value in row["label"]]},
        batched=True,
    )
    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=len(LABELS))
    model = get_peft_model(model, LoraConfig(task_type=TaskType.SEQ_CLS, r=8, lora_alpha=16, lora_dropout=0.1, target_modules=["query", "value"]))
    trainer = Trainer(model=model, args=TrainingArguments(output_dir=str(output_dir), num_train_epochs=2, report_to="none"), train_dataset=tokenized["train"], eval_dataset=tokenized["test"])
    trainer.train()
    trainer.save_model(str(output_dir))
    tokenizer.save_pretrained(str(output_dir))