import os
from transformers import DistilBertForSequenceClassification, Trainer, TrainingArguments
from data import prepare_data
from utils import MyDataset, compute_metrics

def train_model():
    # Task 3: Load Pre-trained model
    train_enc, train_lab, test_enc, test_lab, label2id = prepare_data()
    train_dataset = MyDataset(train_enc, train_lab)
    test_dataset = MyDataset(test_enc, test_lab)

    model = DistilBertForSequenceClassification.from_pretrained(
        'distilbert-base-cased', 
        num_labels=len(label2id)
    )

    # Task 4: W&B Experiment Tracking
    training_args = TrainingArguments(
        output_dir='mlops-assignment2-model',
        num_train_epochs=1,
        per_device_train_batch_size=16,
        logging_dir='./logs',
        logging_steps=10,
        eval_strategy='steps',
        report_to='wandb',  # Enable W&B
        push_to_hub=True,   # Task 5: Hub Deployment
        hub_model_id='mlops-assignment2-distilbert',
        hub_strategy='every_save'
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=test_dataset,
        compute_metrics=compute_metrics
    )
    
    trainer.train()
    trainer.push_to_hub() 
    return trainer, test_dataset

if __name__ == '__main__':
    train_model()