Create train.py

train.py

@@ -0,0 +1,76 @@
+from transformers import TrainingArguments, Trainer
+from datasets import load_dataset
+import evaluate
+import numpy as np
+from modeling_octagon import OctagonForSequenceClassification, OctagonConfig
+from tokenization_octagon import OctagonTokenizer
+
+# Load dataset
+dataset = load_dataset("imdb")
+
+# Sample training (for demo purposes, use smaller subset)
+train_dataset = dataset["train"].shuffle(seed=42).select(range(1000))
+eval_dataset = dataset["test"].shuffle(seed=42).select(range(200))
+
+# Initialize tokenizer
+tokenizer = OctagonTokenizer.train_tokenizer(
+    texts=train_dataset["text"],
+    vocab_size=30522,
+    save_path="octagon-tokenizer.json"
+)
+
+# Tokenize function
+def tokenize_function(examples):
+    return tokenizer(examples["text"], padding="max_length", truncation=True)
+
+tokenized_train = train_dataset.map(tokenize_function, batched=True)
+tokenized_eval = eval_dataset.map(tokenize_function, batched=True)
+
+# Model config
+config = OctagonConfig(
+    vocab_size=30522,
+    hidden_size=128,  # Smaller for demo
+    num_hidden_layers=4,
+    num_attention_heads=4,
+    intermediate_size=512,
+    num_labels=2
+)
+
+model = OctagonForSequenceClassification(config)
+
+# Metrics
+metric = evaluate.load("accuracy")
+
+def compute_metrics(eval_pred):
+    logits, labels = eval_pred
+    predictions = np.argmax(logits, axis=-1)
+    return metric.compute(predictions=predictions, references=labels)
+
+# Training args
+training_args = TrainingArguments(
+    output_dir="octagon_model",
+    evaluation_strategy="epoch",
+    save_strategy="epoch",
+    learning_rate=2e-5,
+    per_device_train_batch_size=8,
+    per_device_eval_batch_size=8,
+    num_train_epochs=3,
+    weight_decay=0.01,
+    load_best_model_at_end=True,
+)
+
+# Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_train,
+    eval_dataset=tokenized_eval,
+    compute_metrics=compute_metrics,
+)
+
+# Train
+trainer.train()
+
+# Save model
+model.save_pretrained("octagon_model")
+tokenizer.save_pretrained("octagon_model")