autotrain set up

.gitignore

@@ -0,0 +1,2 @@
+venv
+results

Dockerfile

@@ -1,2 +1,22 @@
-FROM huggingface/autotrain-advanced:latest
-CMD pip uninstall -y autotrain-advanced && pip install -U autotrain-advanced && autotrain app --host 0.0.0.0 --port 7860 --workers 1
+# Use a Hugging Face image with PyTorch and Transformers
+FROM huggingface/transformers-pytorch-cpu:latest
+
+# Set the working directory inside the container
+WORKDIR /app
+
+# Copy the requirements file first (to leverage Docker's caching mechanism)
+COPY requirements.txt .
+
+# Install and update python3
+RUN apt update
+RUN apt install -y git libsndfile1-dev tesseract-ocr espeak-ng python3 python3-pip ffmpeg
+RUN python3 -m pip install --no-cache-dir --upgrade pip
+
+# Install dependencies
+RUN python3 -m pip install --no-cache-dir -r requirements.txt
+
+# Copy all remaining files into the container
+COPY . .
+
+# Run the training script when the container starts
+CMD ["python3", "train.py"]

requirements.txt

@@ -0,0 +1,5 @@
+transformers>=4.18.0
+datasets>=2.0.0
+torch>=1.10.0
+Pillow>=8.4.0
+accelerate>=0.9.0

train.py

@@ -0,0 +1,56 @@
+from transformers import (
+    AutoModelForImageClassification,
+    AutoImageProcessor,
+    TrainingArguments,
+    Trainer,
+)
+from datasets import load_dataset
+import os
+
+def train():
+    # Load dataset
+    dataset = load_dataset("ylecun/mnist")
+
+    # Load processor and apply preprocessing to the dataset
+    processor = AutoImageProcessor.from_pretrained("SupremoUGH/image-classification-model")
+    
+    def process(examples):
+        images = [img.convert("RGB") for img in examples["image"]]
+        inputs = processor(images=images, return_tensors="pt")
+        inputs["labels"] = examples["label"]
+        return inputs
+    
+    dataset.set_transform(process) # Sometimes `map` instead of `set_transform`
+    
+    # Load model and train it with certain training arguments
+    model = AutoModelForImageClassification.from_pretrained("SupremoUGH/image-classification-model")
+    training_args = TrainingArguments(
+        output_dir="./results",
+        remove_unused_columns=False,  # Preserve input data
+        per_device_train_batch_size=16,  # Reduce batch size for efficiency
+        eval_strategy="steps",
+        num_train_epochs=3,
+        fp16=False,  # Disable fp16 mixed precision
+        save_steps=500,
+        eval_steps=500,
+        logging_steps=100,
+        learning_rate=2e-4,
+        push_to_hub=False,
+    )
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset["train"],
+        eval_dataset=dataset["test"], # Sometimes called "validation"
+    )
+    trainer.train()
+
+    # Save fine-tuned model
+    save_dir = "./saved_model"
+    os.makedirs(save_dir, exist_ok=True)
+    model.save_pretrained(save_dir)
+    print(f"Model saved to {save_dir}")
+
+
+if __name__ == "__main__":
+    train()