Create train.py

train.py

@@ -0,0 +1,40 @@
+from transformers import AutoProcessor, AutoModelForImageClassification, TrainingArguments, Trainer
+from datasets import load_dataset
+import torch
+
+# Load dataset from the 'dataset' folder
+dataset = load_dataset("imagefolder", data_dir="dataset", split="train", label_column="label")
+
+# Load model and processor
+model = AutoModelForImageClassification.from_pretrained("google/siglip2-base-patch16-naflex", num_labels=2)
+processor = AutoProcessor.from_pretrained("google/siglip2-base-patch16-naflex")
+
+# Preprocess the dataset
+def transform(example):
+    inputs = processor(images=example["image"], return_tensors="pt")
+    inputs["label"] = example["label"]
+    return inputs
+
+dataset = dataset.map(transform, batched=True)
+
+# Training setup
+training_args = TrainingArguments(
+    output_dir="./siglip2-meme-classifier",
+    per_device_train_batch_size=8,
+    num_train_epochs=3,
+    save_steps=100,
+    logging_dir="./logs",
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=dataset,
+)
+
+# Start training
+trainer.train()
+
+# Save the fine-tuned model and processor
+model.save_pretrained("model")
+processor.save_pretrained("model")