Updated Model Tutorial Code

README.md

@@ -74,56 +74,41 @@ Use the code below to get started with the model for general finetuning tasks. P
 ```
 import torch
 from datasets import load_dataset, load_metric
-from transformers import AutoTokenizer, AutoModel, TrainingArguments, Trainer
 import evaluate
-metric = evaluate.load("accuracy")
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
+
+# Load the tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("Koodsml/KooBERT")
+model = AutoModelForSequenceClassification.from_pretrained("Koodsml/KooBERT", num_labels=2)
+
 def compute_metrics(eval_pred):
     logits, labels = eval_pred
     predictions = np.argmax(logits, axis=-1)
     return metric.compute(predictions=predictions, references=labels)
 
-# Load the CoLA dataset
-cola_dataset = load_dataset("glue", "cola")
-
-cola_dataset = cola_dataset.rename_column('label', 'labels')
-cola_dataset = cola_dataset.rename_column('sentence', 'text')
-
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("Koodsml/KooBERT")
-model = AutoModel.from_pretrained("Koodsml/KooBERT", num_labels=2)
-
 def tokenize_function(examples):
-    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=128)
+    return tokenizer(examples["text"], padding='max_length', truncation=True, max_length=128)
 
+# Load the CoLA dataset
+dataset = load_dataset("glue","cola")
+dataset = dataset.rename_column('sentence', 'text')
 
-cola_dataset = cola_dataset.map(tokenize_function, batched=True)
+datset_tok = dataset.map(tokenize_function, batched=True)
 
 # Set the device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
 # Define the training arguments
-training_args = TrainingArguments(
-    output_dir='./results',
-    evaluation_strategy='epoch',
-#     eval_steps=100,
-    save_total_limit=1,
-    learning_rate=2e-5,
-    per_device_train_batch_size=8,
-    per_device_eval_batch_size=8,
-    num_train_epochs=3,
-    weight_decay=0.01,
-    push_to_hub=False,
-)
+training_args = TrainingArguments(output_dir="test_trainer", evaluation_strategy="epoch")
 
 # Define the trainer
 trainer = Trainer(
     model=model,
     args=training_args,
-    train_dataset=cola_dataset['train'],
-    eval_dataset=cola_dataset['validation'],
-#     tokenizer=tokenizer,
-    compute_metrics=compute_metrics
+    train_dataset=datset_tok['train'],
+    eval_dataset=datset_tok['validation'],
+    compute_metrics=compute_metrics,
 )
 
 # Fine-tune on the CoLA dataset