Update train.py

train.py

@@ -1,115 +1,111 @@
-import json
-from datasets import Dataset
+#!/usr/bin/env python3
+import torch
 from transformers import (
-    AutoTokenizer,
+    AutoTokenizer, 
     AutoModelForSequenceClassification,
+    TrainingArguments, 
     Trainer,
-    TrainingArguments,
-)
-import evaluate
-import numpy as np
-import os
-from datetime import datetime
-
-# === CONFIG ===
-MODEL_NAME = "bert-base-uncased"
-MODEL_ID = "prelington/acoli"
-DATASET_PATH = "../dataset/test.json"
-OUTPUT_DIR = "./acoli_model"
-LABELS = ["negative", "neutral", "positive"]
-
-# === LOAD DATASET ===
-print("[INFO] Loading dataset from:", DATASET_PATH)
-if not os.path.exists(DATASET_PATH):
-    raise FileNotFoundError(f"Dataset not found at {DATASET_PATH}")
-
-with open(DATASET_PATH, "r", encoding="utf-8") as f:
-    data = json.load(f)
-
-if not isinstance(data, list):
-    raise ValueError("Dataset must be a list of samples!")
-
-dataset = Dataset.from_list(data)
-dataset = dataset.train_test_split(test_size=0.25, seed=42)
-print("[INFO] Dataset loaded successfully!")
-print(dataset)
-
-# === TOKENIZER AND MODEL ===
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(
-    MODEL_NAME, num_labels=len(LABELS)
+    DataCollatorWithPadding
 )
-
-# === PREPROCESSING ===
-def preprocess(batch):
-    enc = tokenizer(
-        batch["text"],
-        truncation=True,
-        padding="max_length",
-        max_length=128,
-    )
-    enc["labels"] = [LABELS.index(label) for label in batch["label"]]
-    return enc
-
-print("[INFO] Tokenizing dataset...")
-tokenized = dataset.map(preprocess, batched=True)
-tokenized.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
-
-# === METRICS ===
-accuracy = evaluate.load("accuracy")
-f1 = 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_score = f1.compute(predictions=predictions, references=labels, average="macro")
-    return {"accuracy": acc["accuracy"], "f1": f1_score["f1"]}
-
-# === TRAINING ARGUMENTS ===
-time_stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
-
-training_args = TrainingArguments(
-    output_dir=f"{OUTPUT_DIR}-{time_stamp}",
-    evaluation_strategy="epoch",
-    save_strategy="epoch",
-    logging_dir=f"./logs_{time_stamp}",
-    logging_strategy="steps",
-    logging_steps=10,
-    num_train_epochs=3,
-    per_device_train_batch_size=4,
-    per_device_eval_batch_size=4,
-    warmup_ratio=0.1,
-    learning_rate=2e-5,
-    weight_decay=0.01,
-    load_best_model_at_end=True,
-    push_to_hub=True,
-    hub_model_id=MODEL_ID,
-    report_to="none",
-)
-
-# === TRAINER ===
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=tokenized["train"],
-    eval_dataset=tokenized["test"],
-    tokenizer=tokenizer,
-    compute_metrics=compute_metrics,
-)
-
-# === TRAIN ===
-print("[INFO] Starting training...")
-trainer.train()
-print("[INFO] Training complete.")
-
-# === SAVE MODEL LOCALLY ===
-os.makedirs(OUTPUT_DIR, exist_ok=True)
-trainer.save_model(OUTPUT_DIR)
-tokenizer.save_pretrained(OUTPUT_DIR)
-print(f"[INFO] Model saved locally to {OUTPUT_DIR}")
-
-# === PUSH TO HUGGING FACE HUB ===
-print("[INFO] Uploading to Hugging Face Hub...")
-trainer.push_to_hub(commit_message="Initial upload of Acoli sentiment model.")
-print(f"[SUCCESS] Model pushed to https://huggingface.co/{MODEL_ID}")
+from datasets import Dataset
+import json
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class AcoliTrainer:
+    def __init__(self, model_name="xlm-roberta-base", num_labels=3):
+        self.model_name = model_name
+        self.num_labels = num_labels
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForSequenceClassification.from_pretrained(
+            model_name, 
+            num_labels=num_labels
+        )
+        
+    def load_data(self, jsonl_path):
+        """Load data from JSONL file"""
+        texts = []
+        labels = []
+        
+        with open(jsonl_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                data = json.loads(line)
+                texts.append(data['text'])
+                labels.append(data['label'])
+                
+        return Dataset.from_dict({
+            'text': texts,
+            'label': labels
+        })
+    
+    def preprocess_function(self, examples):
+        """Tokenize the texts"""
+        return self.tokenizer(
+            examples['text'], 
+            truncation=True, 
+            padding=True,
+            max_length=512
+        )
+    
+    def train(self, train_path, output_dir="./acoli-model"):
+        """Train the model"""
+        
+        # Load and preprocess data
+        logger.info("Loading training data...")
+        dataset = self.load_data(train_path)
+        tokenized_dataset = dataset.map(self.preprocess_function, batched=True)
+        
+        # Split dataset (80% train, 20% validation)
+        train_test_split = tokenized_dataset.train_test_split(test_size=0.2)
+        train_dataset = train_test_split['train']
+        eval_dataset = train_test_split['test']
+        
+        # Training arguments
+        training_args = TrainingArguments(
+            output_dir=output_dir,
+            learning_rate=2e-5,
+            per_device_train_batch_size=8,
+            per_device_eval_batch_size=8,
+            num_train_epochs=3,
+            weight_decay=0.01,
+            evaluation_strategy="epoch",
+            save_strategy="epoch",
+            load_best_model_at_end=True,
+            push_to_hub=False,  # Set to True if you want to push to HF Hub
+        )
+        
+        # Data collator
+        data_collator = DataCollatorWithPadding(tokenizer=self.tokenizer)
+        
+        # Initialize Trainer
+        trainer = Trainer(
+            model=self.model,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            tokenizer=self.tokenizer,
+            data_collator=data_collator,
+        )
+        
+        # Start training
+        logger.info("Starting training...")
+        trainer.train()
+        
+        # Save the model
+        logger.info(f"Saving model to {output_dir}")
+        trainer.save_model(output_dir)
+        self.tokenizer.save_pretrained(output_dir)
+        
+        return trainer
+
+if __name__ == "__main__":
+    # Example usage
+    trainer = AcoliTrainer()
+    
+    # Train the model
+    trained_trainer = trainer.train("path/to/your/data.jsonl")
+    
+    print("Training completed successfully!")