File size: 2,196 Bytes
2466e49 ac0aa1f 2466e49 ac0aa1f 2466e49 ac0aa1f 2466e49 ac0aa1f 2466e49 ac0aa1f 2466e49 ac0aa1f 2466e49 ac0aa1f 2466e49 ac0aa1f 2466e49 ac0aa1f 2466e49 ac0aa1f 2466e49 ac0aa1f 2466e49 ac0aa1f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
import torch
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TrainingArguments,
Trainer,
DataCollatorForLanguageModeling
)
from datasets import load_dataset
import pandas as pd
# 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))
# 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
# 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
)
# Tokenization function
def tokenize_function(examples):
return tokenizer(examples['text'], truncation=True, max_length=1024)
# Prepare dataset
dataset = load_custom_dataset('instructions.csv')
tokenized_dataset = dataset.map(tokenize_function, batched=True)
# 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',
)
# Data Collator
data_collator = DataCollatorForLanguageModeling(
tokenizer=tokenizer,
mlm=False # For causal language modeling
)
# Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_dataset,
data_collator=data_collator,
)
# Start Training
trainer.train()
# Save Model and Tokenizer
trainer.save_model("./ansible-review-model")
tokenizer.save_pretrained("./ansible-review-model")
print("Training Complete!") |