Team Ai
Modelpublic

dnnsdunca/Logical_Algorithm

sourceHugging Faceapache-2.0updated 2y agoView on Hugging Face
0likes
Model_Training.py44 linesDownload Raw Back to root
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')