from transformers import AutoModelForSequenceClassification

model = AutoModelForSequenceClassification.from_pretrained("Sentiment_Classification_DistilBertBase_FT")

from transformers import AutoTokenizer
import torch

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

text = "Layin n bed with a headache ughhhh...waitin o..."
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
outputs = model(**inputs)
logits = outputs.logits

# To get the predicted sentiment label
predicted_label_id = torch.argmax(logits, dim=1).item()

from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer

model = AutoModelForSequenceClassification.from_pretrained(
    "distilbert-base-uncased",
    num_labels=9,  # Number of sentiment labels
    id2label=id2label,  # Mapping from label IDs to label names
    label2id=label2id,  # Mapping from label names to label IDs
)

training_args = TrainingArguments(
    output_dir="Sentiment_Classification_DistilBertBase_FT",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=2,
    weight_decay=0.01,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    push_to_hub=True,  # Push the trained model to the Hugging Face Model Hub
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_ds["train"],  # Your training dataset
    eval_dataset=tokenized_ds["test"],  # Your evaluation dataset
    tokenizer=tokenizer,
    data_collator=data_collator,  # Your data collator function
    compute_metrics=compute_metrics,  # Your metrics function
)