dnnsdunca/Logical_Algorithm
0
1from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments2from datasets import load_dataset3 4# Load dataset - CodeParrot is a good example dataset5dataset = load_dataset('codeparrot/code-to-text')6 7# Load pre-trained model and tokenizer8model = GPT2LMHeadModel.from_pretrained('gpt2-medium')9tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')10 11# Tokenize dataset12def tokenize_function(examples):13 return tokenizer(examples['code'], truncation=True, padding='max_length', max_length=512)14 15tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=['code'])16 17# Training arguments18training_args = TrainingArguments(19 output_dir="./results",20 evaluation_strategy="epoch",21 learning_rate=5e-5,22 per_device_train_batch_size=4,23 per_device_eval_batch_size=4,24 num_train_epochs=3,25 weight_decay=0.01,26 push_to_hub=True,27 hub_model_id='dnnsdunca/UANN',28 hub_token='YOUR_HUGGINGFACE_TOKEN'29)30 31# Trainer32trainer = Trainer(33 model=model,34 args=training_args,35 train_dataset=tokenized_datasets['train'],36 eval_dataset=tokenized_datasets['validation'],37)38 39# Train model40trainer.train()41 42# Save the model43model.save_pretrained('./codegen_model')44tokenizer.save_pretrained('./codegen_model')