Create train.py

train.py

@@ -0,0 +1,104 @@
+import json
+from datasets import Dataset
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    Trainer,
+    TrainingArguments
+)
+import evaluate
+import numpy as np
+import os
+
+# =============== CONFIG =================
+MODEL_NAME = "bert-base-uncased"
+HUB_MODEL_ID = "your-username/acoli"
+DATA_PATH = "../dataset/test.json"
+OUTPUT_DIR = "./acoli_model"
+NUM_EPOCHS = 3
+BATCH_SIZE = 8
+LEARNING_RATE = 2e-5
+SEED = 42
+# ========================================
+
+# 1️⃣ Load dataset
+print("🔹 Loading dataset...")
+with open(DATA_PATH, "r", encoding="utf-8") as f:
+    raw_data = json.load(f)
+dataset = Dataset.from_list(raw_data)
+
+# Split for training and validation
+dataset = dataset.train_test_split(test_size=0.2, seed=SEED)
+
+# 2️⃣ Tokenizer and model
+print("🔹 Loading tokenizer and model...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=3)
+
+# 3️⃣ Tokenization
+def tokenize_function(batch):
+    return tokenizer(batch["text"], padding="max_length", truncation=True)
+
+print("🔹 Tokenizing data...")
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
+tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+tokenized_datasets.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
+
+# 4️⃣ Metrics
+print("🔹 Preparing metrics...")
+accuracy = evaluate.load("accuracy")
+f1_metric = evaluate.load("f1")
+
+def compute_metrics(eval_pred):
+    logits, labels = eval_pred
+    predictions = np.argmax(logits, axis=-1)
+    acc = accuracy.compute(predictions=predictions, references=labels)
+    f1 = f1_metric.compute(predictions=predictions, references=labels, average="macro")
+    return {"accuracy": acc["accuracy"], "f1": f1["f1"]}
+
+# 5️⃣ Training arguments
+print("🔹 Setting up training arguments...")
+training_args = TrainingArguments(
+    output_dir=OUTPUT_DIR,
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    logging_strategy="steps",
+    logging_steps=20,
+    save_total_limit=2,
+    num_train_epochs=NUM_EPOCHS,
+    per_device_train_batch_size=BATCH_SIZE,
+    per_device_eval_batch_size=BATCH_SIZE,
+    learning_rate=LEARNING_RATE,
+    weight_decay=0.01,
+    load_best_model_at_end=True,
+    metric_for_best_model="accuracy",
+    greater_is_better=True,
+    push_to_hub=True,
+    hub_model_id=HUB_MODEL_ID,
+    report_to="none",
+    seed=SEED
+)
+
+# 6️⃣ Trainer setup
+print("🔹 Initializing Trainer...")
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_datasets["train"],
+    eval_dataset=tokenized_datasets["test"],
+    tokenizer=tokenizer,
+    compute_metrics=compute_metrics
+)
+
+# 7️⃣ Start training
+print("🚀 Training started...")
+trainer.train()
+
+# 8️⃣ Save and push to Hugging Face Hub
+print("🔹 Saving model locally...")
+trainer.save_model(OUTPUT_DIR)
+
+print("🔹 Pushing model to Hugging Face Hub...")
+trainer.push_to_hub()
+
+print("✅ Training complete! Model pushed to Hugging Face successfully.")