Update app.py

app.py

@@ -1,63 +1,28 @@
 import gradio as gr
-from datasets import load_dataset
-from transformers import AutoImageProcessor, AutoModelForImageClassification, Trainer, TrainingArguments
+import random
+from PIL import Image
+import time
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 import torch
-import os
+from datasets import load_dataset
 
-# טוענים דאטאסט מהתיקיות
+# טוענים דאטאסט
 dataset = load_dataset("imagefolder", data_dir=".", split={"train": "train[:80%]", "test": "train[80%:]"})
 
-# בוחרים מודל בסיסי
+
+# טוענים מודל בסיס
 checkpoint = "facebook/deit-tiny-patch16-224"
 processor = AutoImageProcessor.from_pretrained(checkpoint)
+
+# במקום להטעין מודל סופי → טוענים ואז משנים את הראש
 model = AutoModelForImageClassification.from_pretrained(
     checkpoint,
     num_labels=3,
     id2label={0: "rock", 1: "paper", 2: "scissors"},
-    label2id={"rock": 0, "paper": 1, "scissors": 2}
-)
-
-# פונקציה לעיבוד התמונות
-def preprocess(examples):
-    images = [x.convert("RGB") for x in examples["image"]]
-    inputs = processor(images=images, return_tensors="pt")
-    inputs["labels"] = examples["label"]
-    return inputs
-
-dataset = dataset.map(preprocess, batched=True)
-
-# הגדרות אימון מהיר
-training_args = TrainingArguments(
-    output_dir="./results",
-    evaluation_strategy="epoch",
-    save_strategy="epoch",
-    per_device_train_batch_size=4,
-    per_device_eval_batch_size=4,
-    num_train_epochs=2,   # ✅ אפוקים מהירים: רק 2!
-    load_best_model_at_end=True,
-    logging_dir='./logs',
-    logging_steps=5,
-)
-
-trainer = Trainer(
-    model=model,
-    args=training_args,
-    train_dataset=dataset["train"],
-    eval_dataset=dataset["test"],
+    label2id={"rock": 0, "paper": 1, "scissors": 2},
+    ignore_mismatched_sizes=True,   # 🔥 קריטי!
 )
 
-# אימון
-trainer.train()
-
-# פונקציה להרצת חיזוי על תמונה חדשה
-def predict(image):
-    inputs = processor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
-    logits = outputs.logits
-    predicted_class_idx = logits.argmax(-1).item()
-    label = model.config.id2label[predicted_class_idx]
-    return label
-
-# בניית אפליקציה
-demo = gr.Interface(fn=predict, inputs="image", outputs="text")
-demo.launch()
+# פונקציית אימון קטנה (נוסיף עוד מעט)
+# פונקציית משחק
+# אפליקציה Gradio