# scripts/train_sentiment_analysis.py
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from datasets import load_dataset

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("CAMeL-Lab/bert-base-camelbert-mix")
model = AutoModelForSequenceClassification.from_pretrained("CAMeL-Lab/bert-base-camelbert-mix", num_labels=3)

# Load and tokenize dataset
dataset = load_dataset("csv", data_files="../data/sentiment_data.csv")

def preprocess_function(examples):
    return tokenizer(examples["text"], truncation=True)

tokenized_dataset = dataset.map(preprocess_function, batched=True)

# Training arguments
training_args = TrainingArguments(
    output_dir="../models/sentiment_analysis_model",
    evaluation_strategy="epoch",
    per_device_train_batch_size=4,
    num_train_epochs=3
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"]
)

trainer.train()