English
shivanircodel commited on
Commit
402354e
·
verified ·
1 Parent(s): 1167ca7

Create train_model.py

Browse files
Files changed (1) hide show
  1. train_model.py +50 -0
train_model.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from datasets import load_dataset
2
+ from transformers import GPT2LMHeadModel, GPT2Tokenizer
3
+ from transformers import Trainer, TrainingArguments
4
+
5
+ # Load dataset (replace with your actual dataset path or URL)
6
+ dataset = load_dataset("path_to_your_dataset")
7
+
8
+ # Load the pre-trained model and tokenizer
9
+ model_name = "gpt2" # Or another model like T5, BART, etc.
10
+ model = GPT2LMHeadModel.from_pretrained(model_name)
11
+ tokenizer = GPT2Tokenizer.from_pretrained(model_name)
12
+
13
+ # Preprocess the dataset: tokenizing the answers and feedback
14
+ def preprocess_function(examples):
15
+ return tokenizer(examples['answer'], truncation=True, padding='max_length', max_length=512)
16
+
17
+ # Apply preprocessing to the dataset
18
+ encoded_dataset = dataset.map(preprocess_function, batched=True)
19
+
20
+ # Define the training arguments
21
+ training_args = TrainingArguments(
22
+ output_dir="./results", # Where to save the model
23
+ evaluation_strategy="epoch", # Evaluate after each epoch
24
+ per_device_train_batch_size=4, # Batch size for training
25
+ per_device_eval_batch_size=4, # Batch size for evaluation
26
+ num_train_epochs=3, # Number of training epochs
27
+ logging_dir='./logs', # Directory for logging
28
+ logging_steps=10,
29
+ save_steps=500, # Save model checkpoint every 500 steps
30
+ warmup_steps=200, # Gradual learning rate warm-up steps
31
+ weight_decay=0.01, # Weight decay for regularization
32
+ load_best_model_at_end=True, # Load the best model based on evaluation metric
33
+ )
34
+
35
+ # Initialize the Trainer
36
+ trainer = Trainer(
37
+ model=model, # The pre-trained model you want to fine-tune
38
+ args=training_args, # The training configuration
39
+ train_dataset=encoded_dataset['train'], # Your training data
40
+ eval_dataset=encoded_dataset['test'], # Your evaluation data
41
+ )
42
+
43
+ # Start training
44
+ trainer.train()
45
+
46
+ # Save the model after training
47
+ model.save_pretrained("./fine_tuned_model")
48
+ tokenizer.save_pretrained("./fine_tuned_model")
49
+
50
+ print("Model fine-tuning complete!")