Update new_review_code.py

new_review_code.py

@@ -1,74 +1,83 @@
-import os
 import torch
-
-from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments,BitsAndBytesConfig
+from transformers import (
+    AutoModelForCausalLM, 
+    AutoTokenizer, 
+    TrainingArguments, 
+    Trainer, 
+    DataCollatorForLanguageModeling
+)
 from datasets import load_dataset
-from trl import SFTTrainer
-from peft import AutoPeftModelForCausalLM, LoraConfig, get_peft_model, prepare_model_for_kbit_training
-from utils import find_all_linear_names, print_trainable_parameters
+import pandas as pd
 
-output_dir="./results"
-model_name ="codellama/CodeLlama-7b-hf"
+# Check GPU availability
+print("CUDA Available:", torch.cuda.is_available())
+print("Current Device:", torch.cuda.current_device())
+print("Device Name:", torch.cuda.get_device_name(0))
 
-dataset = load_dataset('timdettmers/openassistant-guanaco', split="train")
+# Load and prepare dataset
+def load_custom_dataset(file_path):
+    # Read CSV
+    df = pd.read_csv(file_path)
+    
+    # Ensure 'text' column exists
+    if 'text' not in df.columns:
+        raise ValueError("CSV must have a 'text' column")
+    
+    # Convert to Hugging Face dataset
+    dataset = load_dataset('csv', data_files=file_path, split='train')
+    return dataset
 
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16,
-    bnb_4bit_use_double_quant=True,
+# Model and Tokenizer Setup
+model_name = "codellama/CodeLlama-7b-hf"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name, 
+    torch_dtype=torch.float16,  # Use float16 for memory efficiency
+    device_map="auto"  # Automatic device mapping
 )
 
-base_model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, quantization_config=bnb_config)
-base_model.config.use_cache = False
-base_model = prepare_model_for_kbit_training(base_model)
+# Tokenization function
+def tokenize_function(examples):
+    return tokenizer(examples['text'], truncation=True, max_length=1024)
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-tokenizer.pad_token = tokenizer.eos_token
-tokenizer.padding_side = "right"  # Fix weird overflow issue with fp16 training
+# Prepare dataset
+dataset = load_custom_dataset('instructions.csv')
+tokenized_dataset = dataset.map(tokenize_function, batched=True)
 
-# Change the LORA hyperparameters accordingly to fit your use case
-peft_config = LoraConfig(
-    r=32,
-    lora_alpha=16,
-    target_modules=find_all_linear_names(base_model),
-    lora_dropout=0.05,
-    bias="none",
-    task_type="CAUSAL_LM",
+# Training Arguments
+training_args = TrainingArguments(
+    output_dir="./ansible-review-model",
+    overwrite_output_dir=True,
+    num_train_epochs=4,
+    per_device_train_batch_size=2,
+    save_steps=10_000,
+    save_total_limit=2,
+    prediction_loss_only=True,
+    learning_rate=2e-4,
+    warmup_ratio=0.1,
+    fp16=True,  # Use mixed precision
+    logging_dir='./logs',
 )
 
-base_model = get_peft_model(base_model, peft_config)
-print_trainable_parameters(base_model)
-
-# Parameters for training arguments details => https://github.com/huggingface/transformers/blob/main/src/transformers/training_args.py#L158
-training_args = TrainingArguments(
-    per_device_train_batch_size=1,
-    gradient_accumulation_steps=1,
-    gradient_checkpointing =True,
-    max_grad_norm= 0.3,
-    num_train_epochs=3, 
-    learning_rate=1e-4,
-    bf16=True,
-    save_total_limit=3,
-    logging_steps=300,
-    output_dir=output_dir,
-    optim="paged_adamw_32bit",
-    lr_scheduler_type="constant",
-    warmup_ratio=0.05,
+# Data Collator
+data_collator = DataCollatorForLanguageModeling(
+    tokenizer=tokenizer, 
+    mlm=False  # For causal language modeling
 )
 
-trainer = SFTTrainer(
-    base_model,
-    train_dataset=dataset,
-    dataset_text_field="text",
-    tokenizer=tokenizer,
-    max_seq_length=512,
-    args=training_args
+# Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_dataset,
+    data_collator=data_collator,
 )
 
-trainer.train() 
-trainer.save_model(output_dir)
+# Start Training
+trainer.train()
+
+# Save Model and Tokenizer
+trainer.save_model("./ansible-review-model")
+tokenizer.save_pretrained("./ansible-review-model")
 
-output_dir = os.path.join(output_dir, "final_checkpoint")
-trainer.model.save_pretrained(output_dir)
-tokenizer.save_pretrained(output_dir)
+print("Training Complete!")