import pandas as pd
from sklearn.model_selection import train_test_split
from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification, Trainer, TrainingArguments
from datasets import Dataset
import transformers
print("Transformers version:", transformers.__version__)

# Load and encode dataset
df = pd.read_csv("dataset.csv")
label2id = {l: i for i, l in enumerate(df['label'].unique())}
id2label = {i: l for l, i in label2id.items()}
df['label'] = df['label'].map(label2id)

train_texts, val_texts, train_labels, val_labels = train_test_split(df['text'], df['label'], test_size=0.2)

tokenizer = DistilBertTokenizerFast.from_pretrained("distilbert-base-uncased")
train_enc = tokenizer(list(train_texts), truncation=True, padding=True)
val_enc = tokenizer(list(val_texts), truncation=True, padding=True)

train_dataset = Dataset.from_dict({
    'input_ids': train_enc['input_ids'],
    'attention_mask': train_enc['attention_mask'],
    'labels': list(train_labels)
})
val_dataset = Dataset.from_dict({
    'input_ids': val_enc['input_ids'],
    'attention_mask': val_enc['attention_mask'],
    'labels': list(val_labels)
})

model = DistilBertForSequenceClassification.from_pretrained(
    "distilbert-base-uncased", 
    num_labels=len(label2id),
    id2label=id2label,
    label2id=label2id
)

args = TrainingArguments(
    output_dir="./model",
    eval_strategy="epoch", 
    per_device_train_batch_size=8,
    num_train_epochs=4,
    save_total_limit=1,
    logging_dir="./logs"
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset
)

trainer.train()
model.save_pretrained("spam_detector_model")
tokenizer.save_pretrained("spam_detector_model")