wip

Dockerfile

@@ -0,0 +1,35 @@
+# build with: docker build . --tag sctg/roco-idefics3:0.0.5 --tag sctg/roco-idefics3:latest --push
+# run with 
+# docker run --gpus all --user=42420:42420 -e HF_TOKEN=hf_TOKEN -it sctg/roco-idefics3:0.0.5 bash -i /start.sh hf_TOKEN
+FROM nvidia/cuda:11.6.1-devel-ubuntu20.04
+# FROM nvidia/cuda:11.0.3-devel-ubuntu20.04
+# RUN mkdir -p /workspace
+RUN /usr/sbin/addgroup --gid 42420 ovh
+RUN /usr/sbin/useradd -u 42420 --gid 42420 -m -d /workspace -s /bin/bash ovh
+RUN apt update -y && apt-get install -y curl git git-lfs screen
+COPY --chmod=777 start.sh /start.sh
+COPY learn.py /learn.py
+COPY preload.py /preload.py
+# Mandatory to run the jobs in rootless mode
+USER root
+RUN chown -R 42420:42420 /workspace
+USER 42420
+RUN curl -L https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh > /workspace/miniconda.sh
+RUN /bin/bash /workspace/miniconda.sh -b -p /workspace/.miniconda3
+RUN . /workspace/.miniconda3/bin/activate && conda init --all
+RUN . /workspace/.miniconda3/bin/activate \
+    && pip install -U "safetensors>=0.4.5" \
+    && pip install -U "https://github.com/bitsandbytes-foundation/bitsandbytes/releases/download/continuous-release_main/bitsandbytes-0.44.2.dev0-py3-none-manylinux_2_24_x86_64.whl" \
+    && pip install -U git+https://github.com/huggingface/transformers.git\
+    && pip install huggingface_hub[cli] accelerate datasets peft\
+    && pip install -U Pillow \
+    && pip install -U torchvision torchaudio
+RUN . /workspace/.miniconda3/bin/activate && conda install -y jupyter
+# Mandatory to run the jobs in rootless mode
+# USER root
+# RUN chown -R 42420:42420 /workspace
+USER 42420
+WORKDIR /workspace
+# RUN export HOME=/workspace && cd /workspace && . /workspace/.miniconda3/bin/activate \
+#     && mkdir -p /workspace/data \
+#     && python /preload.py

learn.py

@@ -0,0 +1,152 @@
+# Copyright (C) 2024 Ronan Le Meillat
+# License: Apache License 2.0
+# Description: Train the model on the dataset
+import os
+import torch
+
+from huggingface_hub import login as hf_login
+from datasets import load_dataset
+from peft import LoraConfig
+from transformers import AutoProcessor, BitsAndBytesConfig, Idefics3ForConditionalGeneration, TrainingArguments, Trainer
+from datasets.utils.logging import disable_progress_bar
+disable_progress_bar()
+
+HF_TOKEN = ""
+
+if os.environ.get('HF_TOKEN') is not None:
+  HF_TOKEN = os.environ.get('HF_TOKEN')
+  print(f"Hugging Face token found in environment variable")
+
+hf_login(
+  token=HF_TOKEN,
+  add_to_git_credential=True
+)
+dataset_id = "eltorio/ROCO-radiology"
+prompt= "You are an expert radiologist certified with over 15 years of experience in diagnostic imaging, describe this image"
+source_model_id = "HuggingFaceM4/Idefics3-8B-Llama3"
+destination_model_id = "eltorio/IDEFICS3_ROCOv2"
+output_dir = "IDEFICS3_ROCOv2"
+cache_dir = "/workspace/data"
+train_dataset = load_dataset(dataset_id, split="train", cache_dir=cache_dir)
+
+DEVICE = "cuda:0"
+USE_LORA = False
+USE_QLORA = True
+
+processor = AutoProcessor.from_pretrained(
+    source_model_id,
+    do_image_splitting=False
+)
+
+if USE_QLORA or USE_LORA:
+    lora_config = LoraConfig(
+        r=8,
+        lora_alpha=8,
+        lora_dropout=0.1,
+        target_modules='.*(text_model|modality_projection|perceiver_resampler).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$',
+        use_dora=False if USE_QLORA else True,
+        init_lora_weights="gaussian"
+    )
+    if USE_QLORA:
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16
+        )
+    model = Idefics3ForConditionalGeneration.from_pretrained(
+        source_model_id,
+        torch_dtype=torch.float16,
+        quantization_config=bnb_config if USE_QLORA else None,
+    )
+    model.add_adapter(lora_config)
+    model.enable_adapters()
+else:
+    model = Idefics3ForConditionalGeneration.from_pretrained(
+        source_model_id,
+        torch_dtype=torch.float16,
+        _attn_implementation="flash_attention_2", # This works for A100 or H100
+    ).to(DEVICE)
+    
+class MyDataCollator:
+    def __init__(self, processor):
+        self.processor = processor
+        self.image_token_id = processor.tokenizer.additional_special_tokens_ids[
+            processor.tokenizer.additional_special_tokens.index("<image>")
+        ]
+
+    def __call__(self, samples):
+        texts = []
+        images = []
+        for sample in samples:
+            image = sample["image"]
+            answer = sample["caption"]
+            messages = [
+                {
+                    "role": "system",
+                    "content": [
+                        {"type": "text", "text": prompt}
+                    ]
+
+                },
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image"},
+                    ]
+                },
+                {
+                    "role": "assistant",
+                    "content": [
+                        {"type": "text", "text": answer}
+                    ]
+                }
+            ]
+            text = processor.apply_chat_template(messages, add_generation_prompt=False)
+            texts.append(text.strip())
+            images.append([image.convert('RGB')])
+
+        batch = processor(text=texts, images=images, return_tensors="pt", padding=True)
+
+        labels = batch["input_ids"].clone()
+        labels[labels == processor.tokenizer.pad_token_id] = self.image_token_id
+        batch["labels"] = labels
+
+        return batch
+
+data_collator = MyDataCollator(processor)
+
+
+training_args = TrainingArguments(
+    output_dir = output_dir,
+    overwrite_output_dir = False,
+    auto_find_batch_size = True,
+    learning_rate = 2e-4,
+    fp16 = True,
+    per_device_train_batch_size = 2,
+    per_device_eval_batch_size = 2,
+    gradient_accumulation_steps = 8,
+    dataloader_pin_memory = False,
+    save_total_limit = 3,
+    evaluation_strategy = None,
+    save_strategy = "steps",
+    eval_steps = 100,
+    save_steps = 10, # checkpoint each 10 steps
+    resume_from_checkpoint = True,
+    logging_steps = 5,
+    remove_unused_columns = False,
+    push_to_hub = True,
+    label_names = ["labels"],
+    load_best_model_at_end = False,
+    report_to = "none",
+    optim = "paged_adamw_8bit",
+)
+
+trainer = Trainer(
+    model = model,
+    args = training_args,
+    data_collator = data_collator,
+    train_dataset = train_dataset,
+)
+
+trainer.train()
+

start.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+JOB_URL_SCHEME=${JOB_URL_SCHEME:-"http://"}
+JOB_ID=${JOB_ID:-'localhost'}
+JOB_HOST=${JOB_HOST:-'local'}
+export HOME=/workspace
+cd /workspace
+git lfs install
+if [[ -z "$HF_TOKEN" || ! "$HF_TOKEN" =~ ^hf_ ]]; then
+  export HF_TOKEN=$1
+  unset $1
+fi
+
+echo "HF_TOKEN: $HF_TOKEN"
+. /workspace/.bashrc
+. /workspace/.miniconda3/bin/activate
+git clone https://huggingface.co/eltorio/IDEFICS3_ROCOv2
+git config --global credential.helper store
+
+huggingface-cli login --add-to-git-credential --token $HF_TOKEN
+
+
+screen -dmS jupyter bash -c 'jupyter lab --ip=0.0.0.0 --port=8080 --no-browser --allow-root \
+  --notebook-dir=/workspace \
+  --LabApp.token="" \
+  --LabApp.custom_display_url=${JOB_URL_SCHEME}${JOB_ID}-8080.${JOB_HOST} \
+  --LabApp.allow_remote_access=True \
+  --LabApp.allow_origin="*" \
+  --LabApp.disable_check_xsrf=True'
+
+echo "Jupyter Lab is running at ${JOB_URL_SCHEME}${JOB_ID}-8080.${JOB_HOST}"
+exec "$@"