Upload app.py

app.py

@@ -5,9 +5,10 @@ import torchvision
 from PIL import Image
 import numpy as np
 import os
-from huggingface_hub import HfApi, HfFolder, Repository
-from transformers import ViTForImageClassification, Trainer, TrainingArguments
-from datasets import load_dataset
+import shutil
+from huggingface_hub import HfApi, HfFolder
+from transformers import AutoModelForImageClassification, Trainer, TrainingArguments
+from datasets import load_dataset, Dataset
 from sklearn.metrics import accuracy_score
 
 
@@ -19,7 +20,6 @@ def dummy_gpu():
 HF_MODEL = "google/vit-base-patch16-224"
 HF_DATASET = "verytuffcat/recaptcha-dataset"
 HF_REPO = ""
-
 HF_TOKEN = os.getenv("HF_TOKEN", "")
 if os.getenv("HF_REPO"): HF_REPO = os.getenv("HF_REPO")
 if os.getenv("HF_DATASET"): HF_DATASET = os.getenv("HF_DATASET")
@@ -27,32 +27,45 @@ if os.getenv("HF_MODEL"): HF_MODEL = os.getenv("HF_MODEL")
 OUT_DIR = "./new_model"
 
 
+def pil_to_torch(image: Image.Image):
+    return torchvision.transforms.functional.to_tensor(image.convert("RGB").resize((224, 224), Image.BICUBIC))
+
+
 def compute_metrics(eval_pred):
     predictions, labels = eval_pred
     predictions = np.argmax(predictions, axis=1)
-    return dict(accuracy=accuracy_score(predictions, labels))
+    metrics = dict(accuracy=accuracy_score(predictions, labels))
+    return metrics
 
 
 def collate_fn(batch):
-    pixel_values = torch.stack([torchvision.transforms.functional.to_tensor(x["image"].convert("RGB").resize((224, 224), Image.BICUBIC)) for x in batch])
+    pixel_values = torch.stack([pil_to_torch(x["image"]) for x in batch])
     labels = torch.tensor([x["label"] for x in batch])
     return {"pixel_values": pixel_values, "labels": labels}
 
 
-def train(model_id: str, dataset_id: str, repo_id: str, hf_token: str, log_md: str, progress=gr.Progress(track_tqdm=True)):
+def train(model_id: str, dataset_id: str, repo_id: str, hf_token: str, slice: int, log_md: str, progress=gr.Progress(track_tqdm=True)):
     try:
         if not model_id or not dataset_id or not repo_id: raise gr.Error("Fill fields.")
         if not hf_token: hf_token = HF_TOKEN
         if not hf_token: raise gr.Error("Input HF token.")
         HfFolder.save_token(hf_token)
+        api = HfApi(token=hf_token)
+
+        if slice >= 1: dataset = load_dataset(dataset_id, split=f"train[1:{int(slice)+1}]", num_proc=8)
+        else: dataset = load_dataset(dataset_id, split="train", num_proc=8)
+        labels = dataset.features["label"].names
+        label2id, id2label = dict(), dict()
+        for i, label in enumerate(labels):
+            label2id[label] = i
+            id2label[i] = label
 
-        model = ViTForImageClassification.from_pretrained(model_id)
-        dataset = load_dataset(dataset_id, split="train")
+        model = AutoModelForImageClassification.from_pretrained(model_id, label2id=label2id, id2label=id2label, ignore_mismatched_sizes=True)
 
         training_args = TrainingArguments(
             output_dir=OUT_DIR,
             use_cpu=True,
-            no_cuda=True,
+            no_cuda=True, #
             fp16=True,
             optim="adamw_torch",
             lr_scheduler_type="linear",
@@ -61,8 +74,9 @@ def train(model_id: str, dataset_id: str, repo_id: str, hf_token: str, log_md: s
             num_train_epochs=3,
             gradient_accumulation_steps=1,
             use_ipex=True,
+            #eval_strategy="epoch",
             eval_strategy="no",
-            logging_strategy="no",
+            logging_strategy="epoch",
             remove_unused_columns=False,
             push_to_hub=False,
             save_total_limit=2,
@@ -74,15 +88,18 @@ def train(model_id: str, dataset_id: str, repo_id: str, hf_token: str, log_md: s
             data_collator=collate_fn,
             compute_metrics=compute_metrics,
             train_dataset=dataset,
+            #eval_dataset=dataset,
             eval_dataset=None,
         )
-        trainer.train()
+        train_results = trainer.train()
         trainer.save_model(OUT_DIR)
+        trainer.log_metrics("train", train_results.metrics)
+        trainer.save_metrics("train", train_results.metrics)
+        trainer.save_state()
 
-        api = HfApi(token=hf_token)
-        api.create_repo(repo_id=repo_id, private=True, token=hf_token)
-        repo = Repository(local_dir=OUT_DIR, clone_from=repo_id, use_auth_token=hf_token)
-        repo.push_to_hub()
+        api.create_repo(repo_id=repo_id, private=True, exist_ok=True, token=hf_token)
+        api.upload_folder(repo_id=repo_id, folder_path=OUT_DIR, path_in_repo="", token=HF_TOKEN)
+        shutil.rmtree(OUT_DIR)
 
         return log_md
     except Exception as e:
@@ -96,10 +113,12 @@ with gr.Blocks() as demo:
     with gr.Row():
         repo_id = gr.Textbox(label="Output repo", value=HF_REPO, lines=1)
         hf_token = gr.Textbox(label="HF write token", value="", lines=1)
+    with gr.Accordion("Advanced", open=False):
+        slice = gr.Number(label="Slice dataset", info="If 0, use whole dataset", minimum=0, maximum=999999, step=1, value=0)
     train_btn = gr.Button("Train")
     log_md = gr.Markdown(label="Log", value="<br><br>")
 
-    train_btn.click(train, [model_id, dataset_id, repo_id, hf_token, log_md], [log_md])
+    train_btn.click(train, [model_id, dataset_id, repo_id, hf_token, slice, log_md], [log_md])
 
 
 demo.queue().launch()