DatasetDict({
    train: Dataset({
        features: ['review', 'sentiment'],
        num_rows: 36000
    })
    test: Dataset({
        features: ['review', 'sentiment'],
        num_rows: 5000
    })
    eval: Dataset({
        features: ['review', 'sentiment'],
        num_rows: 9000
    })
})

DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"
model_name = 'snoop088/imdb_tuned-bloom1b1-sentiment-classifier'
loaded_model = AutoModelForSequenceClassification.from_pretrained(model_name, 
                                                                  trust_remote_code=True, 
                                                                  num_labels=2,
                                                                  device_map="auto")
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

my_set = pd.read_csv("./data/df_manual.csv")

inputs = tokenizer(list(my_set["review"]), truncation=True, padding="max_length", max_length=256,  return_tensors="pt").to(DEVICE)
outputs = loaded_model(**inputs)
outcome = np.argmax(torch.Tensor.cpu(outputs.logits), axis=-1)

training_arguments = TrainingArguments(
    output_dir="your_tuned_model_name",
    save_strategy="epoch",
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=4,
    optim="adamw_torch",
    evaluation_strategy="steps",
    logging_steps=5,
    learning_rate=1e-5,
    max_grad_norm = 0.3,
    eval_steps=0.2,
    num_train_epochs=2,
    warmup_ratio= 0.1,
    # group_by_length=True,
    fp16=False,
    weight_decay=0.001,
    lr_scheduler_type="constant",
)

peft_model = get_peft_model(model, LoraConfig(
                            task_type="SEQ_CLS",
                            r=16,
                            lora_alpha=16,
                            target_modules=[
                                'query_key_value',
                                'dense'
                            ],
                            bias="none",
                            lora_dropout=0.05, # Conventional
                        ))

def process_data(example):
    item = tokenizer(example["review"], truncation=True, max_length=320) # see if this is OK for dyn padding
    item["labels"] = [ 1 if sent == 'positive' else 0 for sent in example["sentiment"]]
    return item

tokenised_data = tokenised_data.remove_columns(["review", "sentiment"])
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

import evaluate

def compute_metrics(eval_pred):
    # All metrics are already predefined in the HF `evaluate` package
    precision_metric = evaluate.load("precision")
    recall_metric = evaluate.load("recall")
    f1_metric= evaluate.load("f1")
    accuracy_metric = evaluate.load("accuracy")

    logits, labels = eval_pred # eval_pred is the tuple of predictions and labels returned by the model
    predictions = np.argmax(logits, axis=-1)
    precision = precision_metric.compute(predictions=predictions, references=labels)["precision"]
    recall = recall_metric.compute(predictions=predictions, references=labels)["recall"]
    f1 = f1_metric.compute(predictions=predictions, references=labels)["f1"]
    accuracy = accuracy_metric.compute(predictions=predictions, references=labels)["accuracy"]
    # The trainer is expecting a dictionary where the keys are the metrics names and the values are the scores. 
    return {"precision": precision, "recall": recall, "f1-score": f1, 'accuracy': accuracy}