import os
# Configure environment variables BEFORE importing torch
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
os.environ["PYTORCH_NO_CUDA_MEMORY_CACHING"] = "1"

import torch
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    AutoConfig,
    Trainer,
    TrainingArguments,
    DataCollatorWithPadding,
)
from torch.utils.data import Dataset
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
from peft import LoraConfig, get_peft_model
import numpy as np
from transformers.utils import logging

# Set up logging
logging.set_verbosity_info()

# Model and Tokenizer loading
model_name = "meta-llama/Llama-3.2-1B-Instruct"
cache_dir = "/kaggle/input/llama/other/default/1/models/"

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    model_name,
    cache_dir=cache_dir,
    local_files_only=True,
    padding_side="right",
)
tokenizer.pad_token = tokenizer.eos_token

# Load the model configuration
config = AutoConfig.from_pretrained(
    model_name,
    cache_dir=cache_dir,
    local_files_only=True,
)
config.pad_token_id = tokenizer.pad_token_id
config.num_labels = 3

# Check if GPUs are available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.cuda.is_available():
    n_gpus = torch.cuda.device_count()
    print(f"Number of GPUs available: {n_gpus}")
else:
    n_gpus = 0
    print("No GPUs available.")

# Load the model
model = AutoModelForSequenceClassification.from_pretrained(
    model_name,
    config=config,
    cache_dir=cache_dir,
    local_files_only=True,
    device_map={"": device} if device.type == "cuda" else None,  # Load on the specified device
)

# LoRA configuration
lora_config = LoraConfig(
    r=8,
    lora_alpha=16,
    target_modules=["q_proj", "v_proj"],
    lora_dropout=0.05,
    bias="none",
    task_type="SEQ_CLS",
)

# Apply LoRA to the model
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

# Ensure LoRA modules are on the correct devices
def move_lora_to_device(model):
    for name, module in model.named_modules():
        if hasattr(module, 'lora_A') or hasattr(module, 'lora_B'):
            device = next(module.parameters()).device
            if hasattr(module, 'lora_A'):
                module.lora_A.to(device)
            if hasattr(module, 'lora_B'):
                module.lora_B.to(device)

move_lora_to_device(model)

# Data Loading and Preprocessing
df = pd.read_csv("/kaggle/input/amazon-reviews/limited_reviews.csv")

train_texts, val_texts, train_labels, val_labels = train_test_split(
    df['text'].astype(str).tolist(),
    df['sentiment'].tolist(),
    test_size=0.2,
    random_state=42
)

label_map = {'positive': 2, 'negative': 0, 'neutral': 1}
train_labels = [label_map.get(label, -1) for label in train_labels]
val_labels = [label_map.get(label, -1) for label in val_labels]

class SentimentDataset(Dataset):
    def __init__(self, texts, labels):
        self.encodings = tokenizer(
            texts,
            truncation=True,
            padding=True,
            max_length=512,
            return_tensors="pt",
        )
        self.labels = torch.tensor(labels)

    def __getitem__(self, idx):
        item = {key: val[idx] for key, val in self.encodings.items()}
        item['labels'] = self.labels[idx]
        return item

    def __len__(self):
        return len(self.labels)

train_dataset = SentimentDataset(train_texts, train_labels)
val_dataset = SentimentDataset(val_texts, val_labels)

# Data collator
data_collator = DataCollatorWithPadding(tokenizer, padding='longest', return_tensors='pt')

# Define the metrics function
def compute_metrics(pred):
    labels = pred.label_ids
    preds = np.argmax(pred.predictions, axis=-1)
    accuracy = accuracy_score(labels, preds)
    precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='weighted', zero_division=0)
    return {
        'accuracy': accuracy,
        'precision': precision,
        'recall': recall,
        'f1': f1
    }

# Training arguments
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=3,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=16,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=5e-5,
    fp16=True,
    fp16_full_eval=True,
    save_total_limit=2,
    load_best_model_at_end=True,
    logging_steps=50,
    dataloader_num_workers=2,
    report_to="none",
)

# Custom Trainer to handle device placement of labels
from transformers import Trainer

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
        labels = inputs.pop("labels")
        # Move inputs to the correct device
        inputs = {k: v.to(model.device) for k, v in inputs.items()}
        # Forward pass
        outputs = model(**inputs)
        logits = outputs.get("logits")
        # Move labels to the same device as logits
        labels = labels.to(logits.device)
        # Compute loss
        loss_fct = torch.nn.CrossEntropyLoss()
        loss = loss_fct(logits.view(-1, model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss

# Initialize the Custom Trainer
trainer = CustomTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=compute_metrics,
)

# Train the model
trainer.train()

# Evaluate the model
eval_results = trainer.evaluate()
print(f"Evaluation results: {eval_results}")

# Save the model and tokenizer
trainer.save_model("fine_tuned_llama_sentiment_model")
tokenizer.save_pretrained("fine_tuned_llama_sentiment_model")

loading weights file model.safetensors from cache at /kaggle/input/llama/other/default/1/models/models--meta-llama--Llama-3.2-1B-Instruct/snapshots/9213176726f574b556790deb65791e0c5aa438b6/model.safetensors
All model checkpoint weights were used when initializing LlamaForSequenceClassification.

Some weights of LlamaForSequenceClassification were not initialized from the model checkpoint at meta-llama/Llama-3.2-1B-Instruct and are newly initialized: ['score.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
trainable params: 858,112 || all params: 1,236,678,656 || trainable%: 0.0694
/opt/conda/lib/python3.10/site-packages/transformers/training_args.py:1568: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead
  warnings.warn(
PyTorch: setting up devices
/tmp/ipykernel_30/3441435683.py:182: FutureWarning: `tokenizer` is deprecated and will be removed in version 5.0.0 for `CustomTrainer.__init__`. Use `processing_class` instead.
  trainer = CustomTrainer(
You have loaded a model on multiple GPUs. `is_model_parallel` attribute will be force-set to `True` to avoid any unexpected behavior such as device placement mismatching.
Using auto half precision backend
***** Running training *****
  Num examples = 24,000
  Num Epochs = 3
  Instantaneous batch size per device = 2
  Total train batch size (w. parallel, distributed & accumulation) = 32
  Gradient Accumulation steps = 16
  Total optimization steps = 2,250
  Number of trainable parameters = 858,112
/opt/conda/lib/python3.10/multiprocessing/popen_fork.py:66: RuntimeWarning: os.fork() was called. os.fork() is incompatible with multithreaded code, and JAX is multithreaded, so this will likely lead to a deadlock.
  self.pid = os.fork()
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
	- Avoid using `tokenizers` before the fork if possible
	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)
huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...
To disable this warning, you can either:
	- Avoid using `tokenizers` before the fork if possible
	- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)

no need for parallelism since i only have one gpu. fix all warnings and remove parallelism of multi gpus... thanks! Return full code,