Training in progress, step 1000

.gitignore

@@ -0,0 +1,171 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+.git/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+
+# nohup output
+nohup.out
+
+# wandb and output
+wandb/
+output/
+trl/
+

.gitignore~

@@ -0,0 +1,170 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+.git/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+
+# nohup output
+nohup.out
+
+# wandb and output
+wandb/
+output/
+

README.md

@@ -0,0 +1,7 @@
+# Alphapen
+
+This project aims to develop an OCR model for instantaneous text extraction from handwritten documents. The ultimate goal is to seamlessly integrate such a model into computers or mobile phones, allowing for the direct digitalization of handwritten documents using a proprietary pen manufactured by a startup company named [Alphapen](https://alphapen.fr/views/index.html).
+
+# Fine-tuning the TrOCR model
+
+python model.py --log_with wandb --push_to_hub True --hub_model_id hadrakey/alphapen_trocr

adapter_config.json

@@ -0,0 +1,35 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": {
+    "base_model_class": "VisionEncoderDecoderModel",
+    "parent_library": "transformers.models.vision_encoder_decoder.modeling_vision_encoder_decoder"
+  },
+  "base_model_name_or_path": "microsoft/trocr-large-handwritten",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 8,
+  "lora_dropout": 0.1,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 1,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "intermediate.dense",
+    "key",
+    "output.dense",
+    "value",
+    "query"
+  ],
+  "task_type": null,
+  "use_dora": false,
+  "use_rslora": false
+}

adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ced3b228841fe1e8e4710f7320f45cc4083ad8fa319cfdb5d7efe02be2f7159
+size 1811016

config.json

@@ -0,0 +1,177 @@
+{
+  "_name_or_path": "microsoft/trocr-base-handwritten",
+  "architectures": [
+    "VisionEncoderDecoderModel"
+  ],
+  "decoder": {
+    "_name_or_path": "",
+    "activation_dropout": 0.0,
+    "activation_function": "gelu",
+    "add_cross_attention": true,
+    "architectures": null,
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 0,
+    "chunk_size_feed_forward": 0,
+    "classifier_dropout": 0.0,
+    "cross_attention_hidden_size": 768,
+    "d_model": 1024,
+    "decoder_attention_heads": 16,
+    "decoder_ffn_dim": 4096,
+    "decoder_layerdrop": 0.0,
+    "decoder_layers": 12,
+    "decoder_start_token_id": 2,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dropout": 0.1,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 2,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "init_std": 0.02,
+    "is_decoder": true,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layernorm_embedding": true,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 512,
+    "min_length": 0,
+    "model_type": "trocr",
+    "no_repeat_ngram_size": 0,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 1,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "scale_embedding": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": false,
+    "use_learned_position_embeddings": true,
+    "vocab_size": 50265
+  },
+  "decoder_start_token_id": 0,
+  "early_stopping": true,
+  "encoder": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": null,
+    "attention_probs_dropout_prob": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "encoder_stride": 16,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "gelu",
+    "hidden_dropout_prob": 0.0,
+    "hidden_size": 768,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_size": 384,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-12,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "vit",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 12,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_channels": 3,
+    "num_hidden_layers": 12,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 16,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "qkv_bias": false,
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "typical_p": 1.0,
+    "use_bfloat16": false
+  },
+  "eos_token_id": 2,
+  "is_encoder_decoder": true,
+  "length_penalty": 2.0,
+  "max_length": 64,
+  "model_type": "vision-encoder-decoder",
+  "no_repeat_ngram_size": 3,
+  "num_beams": 4,
+  "pad_token_id": 1,
+  "processor_class": "TrOCRProcessor",
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.44.2"
+}

data.py

@@ -0,0 +1,152 @@
+import torch
+from torch.utils.data import Dataset
+from PIL import Image
+import json
+from transformers import TrOCRProcessor
+import pandas as pd
+from sklearn.model_selection import train_test_split
+import glob
+import torchvision.transforms as transforms
+import numpy as np
+
+def prepare_data_frame(root_dir):
+    with open(root_dir) as f:
+        d = json.load(f)
+    filename = [d[i]["word_id"]+ ".png"  for i in range(len(d))]
+    text = [d[i]["text"] for i in range(len(d))]
+    data = {'filename': filename, 'text': text}
+    df = pd.DataFrame(data=data)
+    return df
+
+
+class AphaPenDataset(Dataset):
+    def __init__(self, root_dir, df,  processor, transform=None,  max_target_length=128):
+        self.root_dir = root_dir
+        self.df= df
+        # self.filename, self.text = self.prepare_data()
+        self.processor = processor
+        self.max_target_length = max_target_length
+        self.transform = transform
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        # get file name + text 
+        file_name = self.df.filename[idx]
+        text = self.df.text[idx]
+        # prepare image (i.e. resize + normalize)
+        image = Image.open(self.root_dir + file_name).convert("RGB")
+        if self.transform is not None:
+            image = self.transform(image)
+            img=transforms.ToPILImage()(image)
+            img.save("/mnt/data1/Datasets/AlphaPen/transformed_images/" + file_name)
+        pixel_values = self.processor(image, return_tensors="pt").pixel_values
+        # add labels (input_ids) by encoding the text
+        labels = self.processor.tokenizer(text, 
+                                          padding="max_length", 
+                                          max_length=self.max_target_length).input_ids
+        # important: make sure that PAD tokens are ignored by the loss function
+        labels = [label if label != self.processor.tokenizer.pad_token_id else -100 for label in labels]
+
+        encoding = {"pixel_values": pixel_values.squeeze(), "labels": torch.tensor(labels)}
+        return encoding
+
+    def prepare_data(self):
+        with open(self.path_json) as f:
+            d = json.load(f)
+        filename = [d[i]["image_id"]+ ".png"  for i in range(len(d))]
+        text = [d[i]["text"] for i in range(len(d))]
+        return filename, text
+
+
+class AlphaPenPhi3Dataset(Dataset):
+    def __init__(self, root_dir, dataframe, tokenizer, max_length, image_size):
+        self.dataframe = dataframe
+        self.tokenizer = tokenizer
+        self.tokenizer.padding_side = 'left'
+        self.max_length = max_length
+        self.root_dir = root_dir
+        self.transform = transforms.Compose([
+            transforms.Resize((image_size, image_size)),
+            transforms.ToTensor()
+        ])
+        
+    def __len__(self):
+        return len(self.dataframe)
+
+
+    def __getitem__(self, idx):
+        row = self.dataframe.iloc[idx]
+        text = f"<|user|>\n<|image_1|>What is shown in this image?<|end|><|assistant|>\n {row['text']} <|end|>"
+        image_path = self.root_dir + row['filename']
+        
+        # Tokenize text
+        encodings = self.tokenizer(text, truncation=True, padding='max_length', max_length=self.max_length)
+        
+        try:
+            # Load and transform image
+            image = Image.open(image_path).convert("RGB")
+            image = self.image_transform_function(image)
+        except (FileNotFoundError, IOError):
+            # Skip the sample if the image is not found
+            return None
+        
+        labels = self.tokenizer(row['text'], 
+                                          padding="max_length", 
+                                          max_length=self.max_length).input_ids
+        # important: make sure that PAD tokens are ignored by the loss function
+        labels = [label if label != self.tokenizer.pad_token_id else -100 for label in labels]
+        encodings['pixel_values'] = image
+        encodings['labels'] = labels
+
+        return {key: torch.tensor(val) for key, val in encodings.items()}
+
+
+    def image_transform_function(self, image):
+        image = self.transform(image)
+        return image
+
+
+
+
+if __name__ == "__main__":
+    json_path = "/mnt/data1/Datasets/OCR/Alphapen/label_check/"
+    json_path_b2 = "/mnt/data1/Datasets/OCR/Alphapen/DataBatch2/label_check/"
+    root_dir = "/mnt/data1/Datasets/OCR/Alphapen/clean_data/final_cropped_rotated_"
+    root_dir_b2 = "/mnt/data1/Datasets/OCR/Alphapen/clean_data/final_cropped_rotated_"
+    json_files = glob.glob(json_path + "*.json")
+    json_files_b2 = glob.glob(json_path_b2 + "*.json")
+    root_dir = "/mnt/data1/Datasets/OCR/Alphapen/clean_data/final_cropped_rotated_"
+    df_list_b1 = [prepare_data_frame(file) for file in json_files]
+    df_list_b2 = [prepare_data_frame(file) for file in json_files_b2]
+    # df_list = df_list_b1 + df_list_b2
+    df_b1 = pd.concat(df_list_b1)
+    df_b2 = pd.concat(df_list_b2)
+    
+    df_b1.to_csv("/mnt/data1/Datasets/AlphaPen/" + "testing_data_b1.csv")
+    df_b2.to_csv("/mnt/data1/Datasets/AlphaPen/" + "testing_data_b2.csv")
+    # train_df, test_df = train_test_split(df, test_size=0.15)
+    # # we reset the indices to start from zero
+    # train_df.reset_index(drop=True, inplace=True)
+    # test_df.reset_index(drop=True, inplace=True) 
+    # processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
+    # train_dataset = AphaPenDataset(root_dir=root_dir, df=train_df,  processor=processor)
+    # eval_dataset = AphaPenDataset(root_dir=root_dir, df=test_df,  processor=processor)
+    # print("Number of training examples:", len(train_dataset))
+    # print("Number of validation examples:", len(eval_dataset))
+
+    # encoding = train_dataset[0]
+    # for k,v in encoding.items():
+    #     print(k, v.shape)
+
+    # image = Image.open(train_dataset.root_dir + df.filename[0]).convert("RGB")
+    # print('Label: '+df.text[0])
+    # print(image)
+
+    # labels = encoding['labels']
+    # print(labels)
+
+    # labels[labels == -100] = processor.tokenizer.pad_token_id
+    # label_str = processor.decode(labels, skip_special_tokens=True)
+    # print('Decoded Label:', label_str)

finetune_phi3_vision.py

@@ -0,0 +1,263 @@
+from datasets import Dataset, DatasetDict, Image
+import pandas as pd
+import os
+
+import torch
+from peft import LoraConfig
+from transformers import AutoProcessor, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoModelForVision2Seq
+from datetime import datetime
+import evaluate
+from transformers import TrainingArguments, Trainer, Seq2SeqTrainer, Seq2SeqTrainingArguments
+from sklearn.model_selection import train_test_split
+
+import random
+
+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, examples):
+        texts = []
+        images = []
+        for example in examples:
+            image = example["image"]
+            # print(example["query"])
+            question = example["query"]
+            answer = example["answers"]
+            messages = [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "OCR the text in the image."},
+                        {"type": "image"},
+                        {"type": "text", "text": question}
+                    ]
+                },
+                {
+                    "role": "assistant",
+                    "content": [
+                        {"type": "text", "text": answer}
+                    ]
+                }
+            ]
+            text = processor.apply_chat_template(messages, add_generation_prompt=False)
+            texts.append(text.strip())
+            images.append([image])
+
+        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
+        
+# Define train and test size.
+TRAIN_SAMPLES = 1000
+TEST_SAMPLES = 200
+TEST_SIZE = 0.166 #
+samp_list = [1, 15000, 30000, 45000, 60000, 70000]
+
+# Define the directory containing the images.
+df_path = "/mnt/data1/Datasets/AlphaPen/" + "training_data.csv"
+df = pd.read_csv(df_path)
+df.dropna(inplace=True)
+df["id"] = range(df.shape[0])
+df["query"] = "What is shown in this image?"
+train_df, test_df = train_test_split(df, test_size=0.02, random_state=0)
+
+root_dir = "/mnt/data1/Datasets/OCR/Alphapen/clean_data/final_cropped_rotated_"
+image_paths_train = [root_dir + img for img in train_df.filename]
+image_paths_test = [root_dir + img for img in test_df.filename]
+
+# New batch
+df_path_2 = "/mnt/data1/Datasets/AlphaPen/" + "training_b2.csv"
+df_2 = pd.read_csv(df_path_2)
+df_2.dropna(inplace=True)
+df_2["id"] = range(df_2.shape[0])
+df_2["query"] = "What is shown in this image?"
+train_df_b2, test_df_b2 = train_test_split(df_2, test_size=0.01, random_state=0)
+
+root_dir_2 = "/mnt/data1/Datasets/OCR/Alphapen/DataBatch2/clean_data/cropped_data/cropped_"
+image_paths_2_train = [root_dir_2 + img for img in train_df_b2.filename]
+image_paths_2_test = [root_dir_2 + img for img in test_df_b2.filename]
+
+
+ids_test = range(test_df.shape[0] + test_df_b2.shape[0])
+queries_test = test_df['query'].tolist() + test_df_b2['query'].tolist()
+answers_test = test_df['text'].tolist() + test_df_b2['text'].tolist()
+
+# Create the dataset dictionary.
+
+
+eval_dataset_dict = {
+    'id': ids_test,
+    'image': image_paths_test + image_paths_2_test,
+    'query': queries_test,
+    'answers': answers_test
+}
+
+# Create the dataset.
+
+eval_dataset = Dataset.from_dict(eval_dataset_dict)
+
+# Cast the 'image' column to Image type.
+
+eval_dataset = eval_dataset.cast_column("image", Image())
+
+# Split the dataset into train and test.
+# split_dataset = dataset.train_test_split(test_size=TEST_SIZE, shuffle=False)
+
+# train_dataset = split_dataset["train"]
+# eval_dataset  = split_dataset["test"]
+print(len(eval_dataset))
+# Push the dataset on Hugging Face Hub.
+# split_dataset.push_to_hub("NSTiwari/DocumentIDEFICS_QA")
+
+
+# Define model ID
+# model_id = "microsoft/Phi-3-vision-128k-instruct"
+model_id =  "HuggingFaceM4/idefics2-8b"
+
+DEVICE = "cuda:0"
+USE_LORA = False
+USE_QLORA = True
+
+processor = AutoProcessor.from_pretrained(
+    model_id,
+    do_image_splitting=False
+)
+
+# print(processor.tokenizer.additional_special_tokens.index("<image>"))
+if USE_QLORA or USE_LORA:
+    lora_config = LoraConfig(
+        r=64,
+        lora_alpha=16,
+        lora_dropout=0.1,
+    #     target_modules= [
+    #     "q_proj",
+    #     "k_proj",
+    #     "v_proj",
+    #     "o_proj",
+    #     "gate_proj",
+    #     "up_proj",
+    #     # "down_proj",
+    # ], 
+        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 = AutoModelForVision2Seq.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16,
+        quantization_config=bnb_config if USE_QLORA else None,
+        trust_remote_code=True
+    )
+    model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
+    model.config.pad_token_id = processor.tokenizer.pad_token_id
+    model.config.max_length= 128
+    model.add_adapter(lora_config)
+    model.enable_adapters()
+else:
+    model = AutoModelForVision2Seq.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16,
+        _attn_implementation="flash_attention_2", # Need GPUs like A100 or H100.
+        trust_remote_code=True
+    ).to(DEVICE)
+
+
+
+
+
+data_collator = MyDataCollator(processor)
+
+
+for samp in samp_list:
+    os.environ["WANDB_PROJECT"]="Alphapen"
+    # Create a list of other columns such as id, query, and answer.
+    ids_train = range(train_df.shape[0] + train_df_b2.shape[0])
+    queries_train = train_df['query'].tolist() + train_df_b2['query'].tolist()
+    answers_train = train_df['text'].tolist() + train_df_b2['text'].tolist()
+
+    train_dataset_dict = {
+    'id': ids_train,
+    'image': image_paths_train + image_paths_2_train,
+    'query': queries_train,
+    'answers': answers_train
+    }
+
+    train_dataset = Dataset.from_dict(train_dataset_dict)
+    train_dataset = train_dataset.cast_column("image", Image())
+
+    training_args = Seq2SeqTrainingArguments(
+        predict_with_generate=True,
+        output_dir = "idefics2",
+        learning_rate = 2e-4,
+        fp16 = True,
+        per_device_train_batch_size = 8,
+        per_device_eval_batch_size = 8,
+        gradient_accumulation_steps = 2,
+        dataloader_pin_memory = False,
+        save_total_limit = 3,
+        eval_strategy ="steps",
+        save_strategy = "steps",
+        eval_steps = 500,
+        save_steps = 1000,
+        max_steps = 5000,
+        logging_steps = 10,
+        remove_unused_columns = False,
+        push_to_hub=True,
+        label_names = ["labels"],
+        load_best_model_at_end = False,
+        report_to = "wandb",
+        optim = "paged_adamw_8bit",
+        # run_name=f"idefics2-vision-LoRA-{datetime.now().strftime('%Y-%m-%d-%H-%M-%s')}",
+        run_name="idefics2-vision-LoRA-" + str(samp),
+        hub_model_id="hadrakey/alphapen_idefics2_" + str(samp),
+    )
+
+    def compute_metrics(pred):
+        # accuracy_metric = evaluate.load("precision")
+        cer_metric = evaluate.load("cer")
+
+        labels_ids = pred.label_ids
+        pred_ids = pred.predictions
+        # print(pred_ids)
+        # print(labels_ids)
+        # max_length = max(pred_ids.shape[1], labels_ids.shape[1])
+        # generated_texts = processor.batch_decode(generated_ids[:, inputs["input_ids"].size(1):], skip_special_tokens=True)
+        pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)
+        pred_str = [word.lower() for word in pred_str]
+        # print(pred_str)
+        # pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)
+        labels_ids[labels_ids == -100] = processor.tokenizer.pad_token_id
+        label_str = processor.batch_decode(labels_ids, skip_special_tokens=True)
+        label_str = [word.lower() for word in label_str]
+        # print(label_str)
+        cer = cer_metric.compute(predictions=pred_str, references=label_str)
+        # accuracy = accuracy_metric.compute(predictions=pred_ids.tolist(), references=labels_ids.tolist())
+
+        return {"cer": cer}
+
+
+    trainer = Seq2SeqTrainer(
+        model = model,
+        args = training_args,
+        data_collator = data_collator,
+        train_dataset = train_dataset,
+        eval_dataset = eval_dataset,
+        compute_metrics=compute_metrics,
+    )
+
+    trainer.train()

finetuner_usloath.py

@@ -0,0 +1,174 @@
+# Example inspired from https://huggingface.co/microsoft/Phi-3-vision-128k-instruct
+
+# Import necessary libraries
+from PIL import Image
+import requests
+from transformers import AutoModelForCausalLM
+from transformers import AutoProcessor
+from transformers import BitsAndBytesConfig
+from transformers import  Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator
+import torch
+import pandas as pd
+from torchmetrics.text import CharErrorRate
+from peft import LoraConfig, get_peft_model
+from data import AlphaPenPhi3Dataset
+from sklearn.model_selection import train_test_split
+from datetime import datetime
+import os
+import evaluate
+# tqdm.pandas()
+os.environ["WANDB_PROJECT"]="Alphapen"
+
+# Define model ID
+model_id = "microsoft/Phi-3-vision-128k-instruct"
+# Load data
+
+df_path = "/mnt/data1/Datasets/AlphaPen/" + "training_data.csv"
+df = pd.read_csv(df_path)
+df.dropna(inplace=True)
+train_df, test_df = train_test_split(df, test_size=0.15, random_state=0)
+# we reset the indices to start from zero
+train_df.reset_index(drop=True, inplace=True)
+test_df.reset_index(drop=True, inplace=True)
+root_dir = "/mnt/data1/Datasets/OCR/Alphapen/clean_data/final_cropped_rotated_"
+
+processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+tokenizer = processor.tokenizer
+
+train_dataset = AlphaPenPhi3Dataset(root_dir=root_dir, dataframe=train_df,  tokenizer=tokenizer, max_length=128, image_size=128)
+eval_dataset = AlphaPenPhi3Dataset(root_dir=root_dir, dataframe=test_df.iloc[:10,],  tokenizer=tokenizer, max_length=128, image_size=128)
+
+print(train_dataset[0])
+nf4_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+)
+
+# Load model with 4-bit quantization and map to CUDA
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    trust_remote_code=True,
+    torch_dtype="auto",
+    quantization_config=nf4_config,
+)
+
+# set special tokens used for creating the decoder_input_ids from the labels
+model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
+model.config.pad_token_id = processor.tokenizer.pad_token_id
+# make sure vocab size is set correctly
+# model.config.vocab_size = model.config.decoder.vocab_size
+# for peft
+# model.vocab_size = model.config.decoder.vocab_size
+
+# set beam search parameters
+model.config.eos_token_id = processor.tokenizer.sep_token_id
+model.config.max_new_tokens= 128
+model.config.early_stopping = True
+model.config.no_repeat_ngram_size = 3
+model.config.length_penalty = 2.0
+model.config.num_beams = 4
+
+
+# LoRa
+lora_config = LoraConfig(
+    r=64,
+    lora_alpha=16,
+    lora_dropout=0.1,
+    # target_modules = 'all-linear'
+    target_modules=[
+        "q_proj",
+        "k_proj",
+        "v_proj",
+        "o_proj",
+        # "gate_proj",
+        # "up_proj",
+        # "down_proj",
+    ],
+)
+# print(model)
+# import torch
+# from transformers import Conv1D
+
+# def get_specific_layer_names(model):
+#     # Create a list to store the layer names
+#     layer_names = []
+    
+#     # Recursively visit all modules and submodules
+#     for name, module in model.named_modules():
+#         # Check if the module is an instance of the specified layers
+#         if isinstance(module, (torch.nn.Linear, torch.nn.Embedding, torch.nn.Conv2d, Conv1D)):
+#             # model name parsing 
+
+#             layer_names.append('.'.join(name.split('.')[4:]).split('.')[0])
+    
+#     return layer_names
+
+# print(list(set(get_specific_layer_names(model))))
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# model.to(device)
+
+model = get_peft_model(model, lora_config)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+# print(model.vocab_size)
+# run_name=f"Mistral-7B-SQL-QLoRA-{datetime.now().strftime('%Y-%m-%d-%H-%M-%s')}"
+
+# # Step 3: Define the training arguments
+training_args = Seq2SeqTrainingArguments(
+    predict_with_generate=True,
+    evaluation_strategy="steps",
+    per_device_train_batch_size=8,
+    per_device_eval_batch_size=8,
+    bf16=True,
+    bf16_full_eval=True,
+    output_dir="./",
+    logging_steps=100,
+    save_steps=1000,
+    eval_steps=100,
+    report_to="wandb",
+    run_name=f"phi3-vision-LoRA-{datetime.now().strftime('%Y-%m-%d-%H-%M-%s')}",
+    optim="adamw_torch_fused",
+    lr_scheduler_type="cosine",
+    gradient_accumulation_steps=2,
+    learning_rate=1.0e-4,
+    max_steps=10000,
+    push_to_hub=True,
+    hub_model_id="hadrakey/alphapen_phi3",
+)
+
+def compute_metrics(pred):
+    # accuracy_metric = evaluate.load("precision")
+    cer_metric = evaluate.load("cer")
+
+    labels_ids = pred.label_ids
+    pred_ids = pred.predictions
+    print(labels_ids.shape, pred_ids.shape)
+    max_length = max(pred_ids.shape[1], labels_ids.shape[1])
+
+    pred_str = processor.batch_decode(pred_ids, skip_special_tokens=False,  clean_up_tokenization_spaces=False)
+    print(pred_str)
+    # pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)
+    labels_ids[labels_ids == -100] = tokenizer.pad_token_id
+    label_str = processor.batch_decode(labels_ids, skip_special_tokens=True)
+    print(label_str)
+    cer = cer_metric.compute(predictions=pred_str, references=label_str)
+    # accuracy = accuracy_metric.compute(predictions=pred_ids.tolist(), references=labels_ids.tolist())
+
+    return {"cer": cer}
+
+
+# # Step 5: Define the Trainer
+trainer = Seq2SeqTrainer(
+    model=model,
+    tokenizer=tokenizer,
+    args=training_args,
+    compute_metrics=compute_metrics,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    data_collator=default_data_collator
+)
+
+trainer.train()

idefics2/adapter_config.json

@@ -0,0 +1,26 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "HuggingFaceM4/idefics2-8b",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": false,
+  "init_lora_weights": "gaussian",
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 16,
+  "lora_dropout": 0.1,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 64,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": ".*(text_model|modality_projection|perceiver_resampler).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$",
+  "task_type": null,
+  "use_dora": false,
+  "use_rslora": false
+}

idefics2/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e38855b7b26c79a86d6bc42985348143f602714b85923b6fcf6793830f400de
+size 746528304

idefics2/checkpoint-10000/adapter_config.json

@@ -0,0 +1,26 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "HuggingFaceM4/idefics2-8b",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": false,
+  "init_lora_weights": "gaussian",
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 16,
+  "lora_dropout": 0.1,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 64,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": ".*(text_model|modality_projection|perceiver_resampler).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$",
+  "task_type": null,
+  "use_dora": false,
+  "use_rslora": false
+}

idefics2/checkpoint-10000/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a529961f33fd865061f2d504f10e9dbda5d36ac583ca54c807b178a3eef0a02
+size 746528304

idefics2/checkpoint-10000/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "max_length": 128,
+  "pad_token_id": 0,
+  "transformers_version": "4.42.3"
+}

idefics2/checkpoint-10000/optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e4f18fec102e19e47ad7a4dc5a698a67204bd5a3f9a5e592c8b3c510be2357ad
+size 374548180

idefics2/checkpoint-10000/rng_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4cb5638cdcc03d44751491c48ef9325c702b5d355c5ef610bd485b897821f63
+size 14244

idefics2/checkpoint-10000/scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5351a4d2054ce412bb25cc143ed6232593a8af839dadb43b7080a505b08f3f6
+size 1064

idefics2/checkpoint-10000/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:42f9baac86ca26a549fd61c40400a3efd2e95f6a3486ca7a7482e10ccfbb4ac6
+size 5368

idefics2/checkpoint-8000/adapter_config.json

@@ -0,0 +1,26 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "HuggingFaceM4/idefics2-8b",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": false,
+  "init_lora_weights": "gaussian",
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 16,
+  "lora_dropout": 0.1,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 64,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": ".*(text_model|modality_projection|perceiver_resampler).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$",
+  "task_type": null,
+  "use_dora": false,
+  "use_rslora": false
+}

idefics2/checkpoint-8000/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9beb18aad93c70c527415b17f7e7cfc2d3142ad6521c26cf5b27642f6cfd1d68
+size 746528304

idefics2/checkpoint-8000/generation_config.json

@@ -0,0 +1,18 @@
+{
+  "_from_model_config": true,
+  "bad_words_ids": [
+    [
+      32000
+    ],
+    [
+      32001
+    ]
+  ],
+  "bos_token_id": 1,
+  "eos_token_id": [
+    2,
+    32002
+  ],
+  "pad_token_id": 0,
+  "transformers_version": "4.42.3"
+}

idefics2/checkpoint-8000/optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4b52cc7ac0241e0f119dba0c6b72dc5f4d5a429b38bd94eb75edb5a358b4b644
+size 374548180

idefics2/checkpoint-8000/rng_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca311e524f1174ed34780c49c50e3110c40901c89b57036e41617328bbe51608
+size 14244

idefics2/checkpoint-8000/scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b0b4f2b37c59e9aaf35b1999b2d2c84957e88f5548679884092460d7d6d53d20
+size 1064

idefics2/checkpoint-8000/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81357016d48c35c8aa3b0681c2713c513fb1ee31efc707674d85a48ef9eee341
+size 5368

idefics2/checkpoint-9000/adapter_config.json

@@ -0,0 +1,26 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "HuggingFaceM4/idefics2-8b",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": false,
+  "init_lora_weights": "gaussian",
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 16,
+  "lora_dropout": 0.1,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 64,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": ".*(text_model|modality_projection|perceiver_resampler).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$",
+  "task_type": null,
+  "use_dora": false,
+  "use_rslora": false
+}

idefics2/checkpoint-9000/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:073760f03d80782f132b9c3b74c826df46d8654df042a472117a01290cb7e44f
+size 746528304

idefics2/checkpoint-9000/generation_config.json

@@ -0,0 +1,18 @@
+{
+  "_from_model_config": true,
+  "bad_words_ids": [
+    [
+      32000
+    ],
+    [
+      32001
+    ]
+  ],
+  "bos_token_id": 1,
+  "eos_token_id": [
+    2,
+    32002
+  ],
+  "pad_token_id": 0,
+  "transformers_version": "4.42.3"
+}

idefics2/checkpoint-9000/optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bdd19fe724219197bfb17e4e705cf8801986dcb1e617f05fabedaf8ec38279ee
+size 374548180

idefics2/checkpoint-9000/rng_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a6e481b24fb2e50ba449b0f66b256e83542f9ceba4a6efa543ab9acb0848a1b
+size 14244

idefics2/checkpoint-9000/scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23554f3e4062472629006034b4ce40b01e49bd325d1fb48661fcf4c6868ee807
+size 1064

idefics2/checkpoint-9000/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81357016d48c35c8aa3b0681c2713c513fb1ee31efc707674d85a48ef9eee341
+size 5368

idefics2/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c886ec66a448f0680d0a46cd28b697b6899ecc0627e105de6d1eac26f3c78140
+size 5368

inference.py

@@ -0,0 +1,98 @@
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import pandas as pd
+from PIL import Image
+from torchmetrics.text import CharErrorRate
+
+# Finetuned model
+model_finetune_1 = VisionEncoderDecoderModel.from_pretrained("hadrakey/alphapen_new_large_1")
+model_finetune_2 = VisionEncoderDecoderModel.from_pretrained("hadrakey/alphapen_new_large_15000")
+model_finetune_3 = VisionEncoderDecoderModel.from_pretrained("hadrakey/alphapen_new_large_30000")
+model_finetune_4 = VisionEncoderDecoderModel.from_pretrained("hadrakey/alphapen_new_large_45000")
+model_finetune_5 = VisionEncoderDecoderModel.from_pretrained("hadrakey/alphapen_new_large_60000")
+model_finetune_6 = VisionEncoderDecoderModel.from_pretrained("hadrakey/alphapen_new_large_70000")
+
+#Baseline
+model_base = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
+
+processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
+
+# Checked label
+df_path = "/mnt/data1/Datasets/AlphaPen/" + "testing_data.csv"
+data = pd.read_csv(df_path)
+data.dropna(inplace=True)
+data.reset_index(inplace=True)
+sample = data.iloc[:50,:]
+
+root_dir = "/mnt/data1/Datasets/OCR/Alphapen/clean_data/"
+
+inf_baseline = []
+inf_finetune_1 = []
+inf_finetune_2 = []
+inf_finetune_3 = []
+inf_finetune_4 = []
+inf_finetune_5 = []
+inf_finetune_6 = []
+
+cer_fine_1 = []
+cer_fine_2 = []
+cer_fine_3 = []
+cer_fine_4 = []
+cer_fine_5 = []
+cer_fine_6 = []
+cer_base = []
+
+cer_metric = CharErrorRate()
+
+for idx in range(len(sample)):
+    image = Image.open(root_dir + "final_cropped_rotated_" + sample.filename[idx]).convert("RGB")
+    
+    pixel_values = processor(image, return_tensors="pt").pixel_values
+    generated_ids_base = model_base.generate(pixel_values)
+    generated_ids_fine_1 = model_finetune_1.generate(pixel_values)
+    generated_ids_fine_2= model_finetune_2.generate(pixel_values)
+    generated_ids_fine_3 = model_finetune_3.generate(pixel_values)
+    generated_ids_fine_4 = model_finetune_4.generate(pixel_values)
+    generated_ids_fine_5 = model_finetune_5.generate(pixel_values)
+    generated_ids_fine_6 = model_finetune_6.generate(pixel_values)
+
+    generated_text_base = processor.batch_decode(generated_ids_base, skip_special_tokens=True)[0]
+    generated_text_fine_1= processor.batch_decode(generated_ids_fine_1, skip_special_tokens=True)[0]
+    generated_text_fine_2= processor.batch_decode(generated_ids_fine_2, skip_special_tokens=True)[0]
+    generated_text_fine_3= processor.batch_decode(generated_ids_fine_3, skip_special_tokens=True)[0]
+    generated_text_fine_4= processor.batch_decode(generated_ids_fine_4, skip_special_tokens=True)[0]
+    generated_text_fine_5= processor.batch_decode(generated_ids_fine_5, skip_special_tokens=True)[0]
+    generated_text_fine_6= processor.batch_decode(generated_ids_fine_6, skip_special_tokens=True)[0]
+
+    cer_fine_1.append(cer_metric(generated_text_fine_1.lower(), sample.text[idx].lower()).detach().numpy())
+    cer_fine_2.append(cer_metric(generated_text_fine_2.lower(), sample.text[idx].lower()).detach().numpy())
+    cer_fine_3.append(cer_metric(generated_text_fine_3.lower(), sample.text[idx].lower()).detach().numpy())
+    cer_fine_4.append(cer_metric(generated_text_fine_4.lower(), sample.text[idx].lower()).detach().numpy())
+    cer_fine_5.append(cer_metric(generated_text_fine_5.lower(), sample.text[idx].lower()).detach().numpy())
+    cer_fine_6.append(cer_metric(generated_text_fine_6.lower(), sample.text[idx].lower()).detach().numpy())
+    cer_base.append(cer_metric(generated_text_base.lower(), sample.text[idx].lower()).detach().numpy())
+    
+    inf_baseline.append(generated_text_base)
+    inf_finetune_1.append(generated_text_fine_1)
+    inf_finetune_2.append(generated_text_fine_2)
+    inf_finetune_3.append(generated_text_fine_3)
+    inf_finetune_4.append(generated_text_fine_4)
+    inf_finetune_5.append(generated_text_fine_5)
+    inf_finetune_6.append(generated_text_fine_6)
+
+sample["Baseline"]=inf_baseline
+sample["Finetune_1"]=inf_finetune_1 
+sample["Finetune_2"]=inf_finetune_2  
+sample["Finetune_3"]=inf_finetune_3  
+sample["Finetune_4"]=inf_finetune_4  
+sample["Finetune_5"]=inf_finetune_5   
+sample["Finetune_6"]=inf_finetune_6  
+
+sample["cer_1"]=cer_fine_1 
+sample["cer_2"]=cer_fine_2  
+sample["cer_3"]=cer_fine_3  
+sample["cer_4"]=cer_fine_4  
+sample["cer_5"]=cer_fine_5   
+sample["cer_6"]=cer_fine_6
+sample["cer_base"]=cer_base
+
+sample.to_csv("/mnt/data1/Datasets/AlphaPen/" + "inference_results.csv")

inference_idefics2.py

@@ -0,0 +1,97 @@
+from PIL import Image
+import requests
+from transformers import AutoModelForCausalLM
+from transformers import AutoProcessor
+from transformers import BitsAndBytesConfig
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel, AutoModelForVision2Seq
+import torch
+import pandas as pd
+from torchmetrics.text import CharErrorRate
+from peft import PeftModel, PeftConfig
+from torchmetrics.text import CharErrorRate
+from datasets import Dataset, DatasetDict, Image
+# Define train and test size.
+TRAIN_SAMPLES = 1000
+TEST_SAMPLES = 200
+TEST_SIZE = 0.166 #
+DEVICE = "cuda:0"
+peft_model_id = "hadrakey/alphapen_idefics2_finetune_v1"
+
+config = PeftConfig.from_pretrained(peft_model_id)
+processor = AutoProcessor.from_pretrained(config.base_model_name_or_path, trust_remote_code=True)
+base_model = AutoModelForVision2Seq.from_pretrained(config.base_model_name_or_path, device_map="auto", trust_remote_code=True, torch_dtype="auto")
+model = PeftModel.from_pretrained(base_model, peft_model_id)
+model = model.to(DEVICE)
+
+# Define the directory containing the images.
+df_path = "/mnt/data1/Datasets/AlphaPen/" + "testing_data.csv"
+df = pd.read_csv(df_path)
+df.dropna(inplace=True)
+sample = df.iloc[:5000,:]
+sample.reset_index(inplace=True)
+sample["id"] = range(sample.shape[0])
+sample["query"] = "What is shown in this image?"
+
+root_dir = "/mnt/data1/Datasets/OCR/Alphapen/clean_data/final_cropped_rotated_"
+image_paths = [root_dir + img for img in sample.filename]
+# Create a list of other columns such as id, query, and answer.
+ids = sample['id'].tolist()
+queries = sample['query'].tolist()
+answers = sample['text'].tolist()
+
+# Create the dataset dictionary.
+dataset_dict = {
+    'id': ids,
+    'image': image_paths,
+    'query': queries,
+    'answers': answers
+}
+
+# Create the dataset.
+dataset = Dataset.from_dict(dataset_dict)
+
+# Cast the 'image' column to Image type.
+dataset = dataset.cast_column("image", Image())
+
+# Split the dataset into train and test.
+# split_dataset = dataset.train_test_split(test_size=TEST_SIZE, shuffle=False)
+
+# train_dataset = split_dataset["train"]
+# eval_dataset  = split_dataset["test"]
+
+cer_metric = CharErrorRate()
+cer_idefics = []
+idefics_output = []
+
+for idx in range(len(dataset)):
+
+    test_example = dataset[idx]
+
+    image = test_example["image"]
+    query = test_example["query"]
+
+
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Answer briefly."},
+                {"type": "image"},
+                {"type": "text", "text": query}
+            ]
+        }
+    ]
+
+
+    text = processor.apply_chat_template(messages, add_generation_prompt=True)
+    inputs = processor(text=[text.strip()], images=[image], return_tensors="pt", padding=True)
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+    generated_ids = model.generate(**inputs, max_new_tokens=64)
+    generated_texts = processor.batch_decode(generated_ids[:, inputs["input_ids"].size(1):], skip_special_tokens=True)
+    idefics_output.append(generated_texts[0])
+    cer_idefics.append(cer_metric(generated_texts[0].lower(), test_example["answers"].lower()).detach().numpy())
+    # print(generated_texts, test_example["answers"], cer_idefics)
+
+sample["idefics"] = idefics_output
+sample["cer"] = cer_idefics
+sample.to_csv("/mnt/data1/Datasets/AlphaPen/" + "sample_idefics_v1.csv")

model.py

@@ -0,0 +1,204 @@
+
+from dataclasses import dataclass, field
+from typing import Optional
+import pandas as pd
+
+import torch
+from accelerate import Accelerator
+from datasets import load_dataset, Dataset, load_metric
+from peft import LoraConfig
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig, HfArgumentParser, TrainingArguments, VisionEncoderDecoderModel, TrOCRProcessor, Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator, EarlyStoppingCallback
+
+
+from trl import SFTTrainer, is_xpu_available
+from data import AphaPenDataset
+import evaluate
+from sklearn.model_selection import train_test_split
+
+import torchvision.transforms as transforms
+# from utils import compute_metrics 
+from src.calibrator import EncoderDecoderCalibrator
+from src.loss import MarginLoss, KLRegularization
+from src.similarity import CERSimilarity
+import os
+tqdm.pandas()
+
+os.environ["WANDB_PROJECT"]="Alphapen"
+# Define and parse arguments.
+@dataclass
+class ScriptArguments:
+    """
+    The name of the OCR model we wish to fine with Seq2SeqTrainer
+    """
+
+    model_name: Optional[str] = field(default="microsoft/trocr-base-handwritten", metadata={"help": "the model name"})
+    dataset_name: Optional[str] = field(
+        default="Anthropic/hh-rlhf", metadata={"help": "the dataset name"}
+    )
+    log_with: Optional[str] = field(default="none", metadata={"help": "use 'wandb' to log with wandb"})
+    learning_rate: Optional[float] = field(default=1.41e-5, metadata={"help": "the learning rate"})
+    batch_size: Optional[int] = field(default=8, metadata={"help": "the batch size"})
+    seq_length: Optional[int] = field(default=512, metadata={"help": "Input sequence length"})
+    gradient_accumulation_steps: Optional[int] = field(
+        default=16, metadata={"help": "the number of gradient accumulation steps"}
+    )
+    load_in_8bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 8 bits precision"})
+    load_in_4bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 4 bits precision"})
+    use_peft: Optional[bool] = field(default=False, metadata={"help": "Wether to use PEFT or not to train adapters"})
+    trust_remote_code: Optional[bool] = field(default=False, metadata={"help": "Enable `trust_remote_code`"})
+    output_dir: Optional[str] = field(default="output", metadata={"help": "the output directory"})
+    peft_lora_r: Optional[int] = field(default=64, metadata={"help": "the r parameter of the LoRA adapters"})
+    peft_lora_alpha: Optional[int] = field(default=16, metadata={"help": "the alpha parameter of the LoRA adapters"})
+    logging_steps: Optional[int] = field(default=1, metadata={"help": "the number of logging steps"})
+    use_auth_token: Optional[bool] = field(default=True, metadata={"help": "Use HF auth token to access the model"})
+    num_train_epochs: Optional[int] = field(default=3, metadata={"help": "the number of training epochs"})
+    max_steps: Optional[int] = field(default=-1, metadata={"help": "the number of training steps"})
+    max_length: Optional[int] = field(default=10, metadata={"help": "the maximum length"})
+    no_repeat_ngram_size: Optional[int] = field(default=3, metadata={"help": "the number of repeat"})
+    length_penalty: Optional[float] = field(default=2.0, metadata={"help": "the length of penalty"})
+    num_beams: Optional[int] = field(default=3, metadata={"help": "the number of beam search"})
+    early_stopping: Optional[bool] = field(default=True, metadata={"help": "Early stopping"})
+    save_steps: Optional[int] = field(
+        default=1000, metadata={"help": "Number of updates steps before two checkpoint saves"}
+    )
+    save_total_limit: Optional[int] = field(default=10, metadata={"help": "Limits total number of checkpoints."})
+    push_to_hub: Optional[bool] = field(default=False, metadata={"help": "Push the model to HF Hub"})
+    gradient_checkpointing: Optional[bool] = field(
+        default=False, metadata={"help": "Whether to use gradient checkpointing or no"}
+    )
+    gradient_checkpointing_kwargs: Optional[dict] = field(
+        default=None,
+        metadata={
+            "help": "key word arguments to be passed along `torch.utils.checkpoint.checkpoint` method - e.g. `use_reentrant=False`"
+        },
+    )
+    hub_model_id: Optional[str] = field(default=None, metadata={"help": "The name of the model on HF Hub"})
+
+parser = HfArgumentParser(ScriptArguments)
+script_args = parser.parse_args_into_dataclasses()[0]
+
+# # Step 1: Load the dataset
+df_path = "/mnt/data1/Datasets/AlphaPen/" + "training_data.csv"
+df = pd.read_csv(df_path)
+df.dropna(inplace=True)
+train_df, test_df = train_test_split(df, test_size=0.15, random_state=0)
+# we reset the indices to start from zero
+train_df.reset_index(drop=True, inplace=True)
+test_df.reset_index(drop=True, inplace=True)
+root_dir = "/mnt/data1/Datasets/OCR/Alphapen/clean_data/final_cropped_rotated_"
+processor = TrOCRProcessor.from_pretrained(script_args.model_name)
+  
+
+
+train_dataset = AphaPenDataset(root_dir=root_dir, df=train_df,  processor=processor)
+eval_dataset = AphaPenDataset(root_dir=root_dir, df=test_df,  processor=processor)
+
+# Step 2: Load the model
+if script_args.load_in_8bit and script_args.load_in_4bit:
+    raise ValueError("You can't load the model in 8 bits and 4 bits at the same time")
+elif script_args.load_in_8bit or script_args.load_in_4bit:
+    quantization_config = BitsAndBytesConfig(
+        load_in_8bit=script_args.load_in_8bit, load_in_4bit=script_args.load_in_4bit
+    )
+    # Copy the model to each device
+    device_map = (
+        {"": f"xpu:{Accelerator().local_process_index}"}
+        if is_xpu_available()
+        else {"": Accelerator().local_process_index}
+    )
+    torch_dtype = torch.bfloat16
+else:
+    device_map = None
+    quantization_config = None
+    torch_dtype = None
+
+model = VisionEncoderDecoderModel.from_pretrained(
+    script_args.model_name,
+    quantization_config=quantization_config,
+    device_map=device_map,
+    trust_remote_code=script_args.trust_remote_code,
+    torch_dtype=torch_dtype,
+    token=script_args.use_auth_token,
+)
+
+# set special tokens used for creating the decoder_input_ids from the labels
+model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
+model.config.pad_token_id = processor.tokenizer.pad_token_id
+# make sure vocab size is set correctly
+model.config.vocab_size = model.config.decoder.vocab_size
+
+# set beam search parameters
+model.config.eos_token_id = processor.tokenizer.sep_token_id
+model.config.max_length = script_args.max_length
+model.config.early_stopping = script_args.early_stopping
+model.config.no_repeat_ngram_size = script_args.no_repeat_ngram_size
+model.config.length_penalty = script_args.length_penalty
+model.config.num_beams = script_args.num_beams
+
+tokenizer = processor.tokenizer
+sim = CERSimilarity(tokenizer)
+loss = MarginLoss(sim, beta=0.1, num_samples=60)
+reg = KLRegularization(model)
+calibrator = EncoderDecoderCalibrator(model, loss, reg, 15, 15)
+
+
+# # Step 3: Define the training arguments
+training_args = Seq2SeqTrainingArguments(
+    predict_with_generate=True,
+    evaluation_strategy="steps",
+    per_device_train_batch_size=script_args.batch_size,
+    per_device_eval_batch_size=script_args.batch_size,
+    fp16=True, 
+    output_dir=script_args.output_dir,
+    logging_steps=script_args.logging_steps,
+    save_steps=script_args.save_steps,
+    eval_steps=100,
+    save_total_limit=script_args.save_total_limit,
+    # load_best_model_at_end = True,
+    report_to=script_args.log_with,
+    num_train_epochs=script_args.num_train_epochs,
+    push_to_hub=script_args.push_to_hub,
+    hub_model_id=script_args.hub_model_id,
+    gradient_checkpointing=script_args.gradient_checkpointing,
+    # metric_for_best_model="eval/cer"
+    # TODO: uncomment that on the next release 
+    # gradient_checkpointing_kwargs=script_args.gradient_checkpointing_kwargs,
+)
+
+
+# Step 4: Define a metric
+
+def compute_metrics(pred):
+    # accuracy_metric = evaluate.load("precision")
+    cer_metric = evaluate.load("cer")
+
+    labels_ids = pred.label_ids
+    pred_ids = pred.predictions
+
+    pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)
+    labels_ids[labels_ids == -100] = processor.tokenizer.pad_token_id
+    label_str = processor.batch_decode(labels_ids, skip_special_tokens=True)
+
+    cer = cer_metric.compute(predictions=pred_str, references=label_str)
+    # accuracy = accuracy_metric.compute(predictions=pred_ids.tolist(), references=labels_ids.tolist())
+
+    return {"cer": cer}
+
+early_stop = EarlyStoppingCallback(10, .001)
+# # Step 5: Define the Trainer
+trainer = Seq2SeqTrainer(
+    model=model,
+    tokenizer=processor.feature_extractor,
+    args=training_args,
+    compute_metrics=compute_metrics,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    data_collator=default_data_collator,
+    # callbacks = [early_stop]
+)
+
+trainer.train()
+
+# # Step 6: Save the model
+# trainer.save_model(script_args.output_dir)

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b14472ca382e9d96ea7efd3c778cbf0b73a412e31bc41cfec8d97e8988e6063d
+size 1335747032

model_sft.py

@@ -0,0 +1,217 @@
+
+from dataclasses import dataclass, field
+from typing import Optional
+import pandas as pd
+
+import torch
+from accelerate import Accelerator
+from datasets import load_dataset, Dataset, load_metric
+from peft import LoraConfig
+from tqdm import tqdm
+from transformers import AutoModelForCausalLM, BitsAndBytesConfig, HfArgumentParser, TrainingArguments, VisionEncoderDecoderModel, TrOCRProcessor, Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator, EarlyStoppingCallback
+
+
+from trl import SFTTrainer, is_xpu_available
+from data import AphaPenDataset
+import evaluate
+from sklearn.model_selection import train_test_split
+
+import torchvision.transforms as transforms
+# from utils import compute_metrics 
+
+tqdm.pandas()
+
+
+# Define and parse arguments.
+@dataclass
+class ScriptArguments:
+    """
+    The name of the OCR model we wish to fine with Seq2SeqTrainer
+    """
+
+    model_name: Optional[str] = field(default="microsoft/trocr-base-handwritten", metadata={"help": "the model name"})
+    dataset_name: Optional[str] = field(
+        default="Anthropic/hh-rlhf", metadata={"help": "the dataset name"}
+    )
+    log_with: Optional[str] = field(default="none", metadata={"help": "use 'wandb' to log with wandb"})
+    learning_rate: Optional[float] = field(default=1.41e-5, metadata={"help": "the learning rate"})
+    batch_size: Optional[int] = field(default=8, metadata={"help": "the batch size"})
+    seq_length: Optional[int] = field(default=512, metadata={"help": "Input sequence length"})
+    gradient_accumulation_steps: Optional[int] = field(
+        default=16, metadata={"help": "the number of gradient accumulation steps"}
+    )
+    load_in_8bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 8 bits precision"})
+    load_in_4bit: Optional[bool] = field(default=False, metadata={"help": "load the model in 4 bits precision"})
+    use_peft: Optional[bool] = field(default=False, metadata={"help": "Wether to use PEFT or not to train adapters"})
+    trust_remote_code: Optional[bool] = field(default=False, metadata={"help": "Enable `trust_remote_code`"})
+    output_dir: Optional[str] = field(default="output", metadata={"help": "the output directory"})
+    peft_lora_r: Optional[int] = field(default=64, metadata={"help": "the r parameter of the LoRA adapters"})
+    peft_lora_alpha: Optional[int] = field(default=16, metadata={"help": "the alpha parameter of the LoRA adapters"})
+    logging_steps: Optional[int] = field(default=1, metadata={"help": "the number of logging steps"})
+    use_auth_token: Optional[bool] = field(default=True, metadata={"help": "Use HF auth token to access the model"})
+    num_train_epochs: Optional[int] = field(default=3, metadata={"help": "the number of training epochs"})
+    max_steps: Optional[int] = field(default=-1, metadata={"help": "the number of training steps"})
+    max_length: Optional[int] = field(default=10, metadata={"help": "the maximum length"})
+    no_repeat_ngram_size: Optional[int] = field(default=3, metadata={"help": "the number of repeat"})
+    length_penalty: Optional[float] = field(default=2.0, metadata={"help": "the length of penalty"})
+    num_beams: Optional[int] = field(default=3, metadata={"help": "the number of beam search"})
+    early_stopping: Optional[bool] = field(default=True, metadata={"help": "Early stopping"})
+    save_steps: Optional[int] = field(
+        default=1000, metadata={"help": "Number of updates steps before two checkpoint saves"}
+    )
+    save_total_limit: Optional[int] = field(default=10, metadata={"help": "Limits total number of checkpoints."})
+    push_to_hub: Optional[bool] = field(default=False, metadata={"help": "Push the model to HF Hub"})
+    gradient_checkpointing: Optional[bool] = field(
+        default=False, metadata={"help": "Whether to use gradient checkpointing or no"}
+    )
+    gradient_checkpointing_kwargs: Optional[dict] = field(
+        default=None,
+        metadata={
+            "help": "key word arguments to be passed along `torch.utils.checkpoint.checkpoint` method - e.g. `use_reentrant=False`"
+        },
+    )
+    hub_model_id: Optional[str] = field(default=None, metadata={"help": "The name of the model on HF Hub"})
+
+parser = HfArgumentParser(ScriptArguments)
+script_args = parser.parse_args_into_dataclasses()[0]
+
+# # Step 1: Load the dataset
+df_path = "/mnt/data1/Datasets/AlphaPen/" + "training_data.csv"
+df = pd.read_csv(df_path)
+df.dropna(inplace=True)
+train_df, test_df = train_test_split(df, test_size=0.15, random_state=0)
+# we reset the indices to start from zero
+train_df.reset_index(drop=True, inplace=True)
+test_df.reset_index(drop=True, inplace=True)
+root_dir = "/mnt/data1/Datasets/OCR/Alphapen/clean_data/final_cropped_rotated_"
+processor = TrOCRProcessor.from_pretrained(script_args.model_name)
+
+ # Transformation for training including augmentations
+transform = transforms.Compose([ 
+    transforms.PILToTensor(),
+    transforms.RandomRotation(degrees=(0, 180)) 
+]) 
+  
+
+
+train_dataset = AphaPenDataset(root_dir=root_dir, df=train_df,  processor=processor, transform=transform)
+eval_dataset = AphaPenDataset(root_dir=root_dir, df=test_df,  processor=processor)
+
+# Step 2: Load the model
+if script_args.load_in_8bit and script_args.load_in_4bit:
+    raise ValueError("You can't load the model in 8 bits and 4 bits at the same time")
+elif script_args.load_in_8bit or script_args.load_in_4bit:
+    quantization_config = BitsAndBytesConfig(
+        load_in_8bit=script_args.load_in_8bit, load_in_4bit=script_args.load_in_4bit
+    )
+    # Copy the model to each device
+    device_map = (
+        {"": f"xpu:{Accelerator().local_process_index}"}
+        if is_xpu_available()
+        else {"": Accelerator().local_process_index}
+    )
+    torch_dtype = torch.bfloat16
+else:
+    device_map = None
+    quantization_config = None
+    torch_dtype = None
+
+model = VisionEncoderDecoderModel.from_pretrained(
+    script_args.model_name,
+    quantization_config=quantization_config,
+    device_map=device_map,
+    trust_remote_code=script_args.trust_remote_code,
+    torch_dtype=torch_dtype,
+    token=script_args.use_auth_token,
+)
+
+# set special tokens used for creating the decoder_input_ids from the labels
+model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
+model.config.pad_token_id = processor.tokenizer.pad_token_id
+# make sure vocab size is set correctly
+model.config.vocab_size = model.config.decoder.vocab_size
+
+# set beam search parameters
+model.config.eos_token_id = processor.tokenizer.sep_token_id
+model.config.max_length = script_args.max_length
+model.config.early_stopping = script_args.early_stopping
+model.config.no_repeat_ngram_size = script_args.no_repeat_ngram_size
+model.config.length_penalty = script_args.length_penalty
+model.config.num_beams = script_args.num_beams
+
+
+
+
+# # Step 3: Define the training arguments
+training_args = Seq2SeqTrainingArguments(
+    predict_with_generate=True,
+    evaluation_strategy="steps",
+    # per_device_train_batch_size=script_args.batch_size,
+    # per_device_eval_batch_size=script_args.batch_size,
+    fp16=True, 
+    output_dir=script_args.output_dir,
+    logging_steps=script_args.logging_steps,
+    save_steps=script_args.save_steps,
+    eval_steps=100,
+    save_total_limit=script_args.save_total_limit,
+    load_best_model_at_end = True,
+    report_to=script_args.log_with,
+    num_train_epochs=script_args.num_train_epochs,
+    push_to_hub=script_args.push_to_hub,
+    hub_model_id=script_args.hub_model_id,
+    gradient_checkpointing=script_args.gradient_checkpointing,
+    auto_find_batch_size=True,
+    metric_for_best_model="eval/cer"
+    # TODO: uncomment that on the next release 
+    # gradient_checkpointing_kwargs=script_args.gradient_checkpointing_kwargs,
+)
+
+
+# Step 4: Define a metric
+
+def compute_metrics(pred):
+    # accuracy_metric = evaluate.load("precision")
+    cer_metric = evaluate.load("cer")
+
+    labels_ids = pred.label_ids
+    pred_ids = pred.predictions
+
+    pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)
+    labels_ids[labels_ids == -100] = processor.tokenizer.pad_token_id
+    label_str = processor.batch_decode(labels_ids, skip_special_tokens=True)
+
+    cer = cer_metric.compute(predictions=pred_str, references=label_str)
+    # accuracy = accuracy_metric.compute(predictions=pred_ids.tolist(), references=labels_ids.tolist())
+
+    return {"cer": cer}
+
+early_stop = EarlyStoppingCallback(10, .001)
+
+# Step 5: Define the LoraConfig
+if script_args.use_peft:
+    peft_config = LoraConfig(
+        r=script_args.peft_lora_r,
+        lora_alpha=script_args.peft_lora_alpha,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules="all-linear"
+    )
+else:
+    peft_config = None
+# # Step 6: Define the Trainer
+trainer = SFTTrainer(
+    model=model,
+    tokenizer=processor.feature_extractor,
+    args=training_args,
+    compute_metrics=compute_metrics,
+    train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
+    data_collator=default_data_collator,
+    peft_config=peft_config,
+    callbacks=[EarlyStoppingCallback(early_stopping_patience=10)]
+)
+
+trainer.train()
+
+# # Step 6: Save the model
+# trainer.save_model(script_args.output_dir)

phi3/checkpoint-25/adapter_config.json

@@ -0,0 +1,26 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "HuggingFaceM4/idefics2-8b",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": false,
+  "init_lora_weights": "gaussian",
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 16,
+  "lora_dropout": 0.1,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 64,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": ".*(text_model|modality_projection|perceiver_resampler).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$",
+  "task_type": null,
+  "use_dora": false,
+  "use_rslora": false
+}

phi3/checkpoint-25/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c804cede36291fb0feb5cee74f4ffeeaec1178864af130c873b410c6f1fe1a18
+size 746528304

phi3/checkpoint-25/generation_config.json

@@ -0,0 +1,18 @@
+{
+  "_from_model_config": true,
+  "bad_words_ids": [
+    [
+      32000
+    ],
+    [
+      32001
+    ]
+  ],
+  "bos_token_id": 1,
+  "eos_token_id": [
+    2,
+    32002
+  ],
+  "pad_token_id": 0,
+  "transformers_version": "4.42.3"
+}

phi3/checkpoint-25/optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d93bc5e374a2ab92fce0a68bbe1800baf7f1ea49f5528c1bd22a6bb987d7a79
+size 374547732

phi3/checkpoint-25/rng_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3bc84e110eb9a1be206e02c97bf5c5d7133033f306401f2e818d8847834cab9f
+size 14244

phi3/checkpoint-25/scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:630fe59a784956405be1a950d9ce52e5bf6a2f1c12f3a8bd4f3869766a5850cd
+size 1064

phi3/checkpoint-25/trainer_state.json

@@ -0,0 +1,84 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.008670019074041963,
+  "eval_steps": 10,
+  "global_step": 25,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0017340038148083927,
+      "grad_norm": 11.838030815124512,
+      "learning_rate": 0.00017600000000000002,
+      "loss": 12.8857,
+      "step": 5
+    },
+    {
+      "epoch": 0.0034680076296167853,
+      "grad_norm": 0.8977920413017273,
+      "learning_rate": 0.00013600000000000003,
+      "loss": 0.6798,
+      "step": 10
+    },
+    {
+      "epoch": 0.0034680076296167853,
+      "eval_loss": 0.2337629497051239,
+      "eval_runtime": 675.29,
+      "eval_samples_per_second": 13.597,
+      "eval_steps_per_second": 1.7,
+      "step": 10
+    },
+    {
+      "epoch": 0.0052020114444251775,
+      "grad_norm": 0.34665364027023315,
+      "learning_rate": 9.6e-05,
+      "loss": 0.1571,
+      "step": 15
+    },
+    {
+      "epoch": 0.006936015259233571,
+      "grad_norm": 0.26853781938552856,
+      "learning_rate": 5.6000000000000006e-05,
+      "loss": 0.1088,
+      "step": 20
+    },
+    {
+      "epoch": 0.006936015259233571,
+      "eval_loss": 0.09983003884553909,
+      "eval_runtime": 686.1959,
+      "eval_samples_per_second": 13.381,
+      "eval_steps_per_second": 1.673,
+      "step": 20
+    },
+    {
+      "epoch": 0.008670019074041963,
+      "grad_norm": 0.22648762166500092,
+      "learning_rate": 1.6000000000000003e-05,
+      "loss": 0.0911,
+      "step": 25
+    }
+  ],
+  "logging_steps": 5,
+  "max_steps": 25,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 25,
+  "stateful_callbacks": {
+    "TrainerControl": {
+      "args": {
+        "should_epoch_stop": false,
+        "should_evaluate": false,
+        "should_log": false,
+        "should_save": true,
+        "should_training_stop": true
+      },
+      "attributes": {}
+    }
+  },
+  "total_flos": 1803863725327872.0,
+  "train_batch_size": 8,
+  "trial_name": null,
+  "trial_params": null
+}

phi3/checkpoint-25/training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f453fa98240e7ba800f54183433d2dfc198cd72c7267c34a5547639a1d49da5c
+size 5112

phi3_ocr.py

@@ -0,0 +1,176 @@
+# Example inspired from https://huggingface.co/microsoft/Phi-3-vision-128k-instruct
+
+# Import necessary libraries
+from PIL import Image
+import requests
+from transformers import AutoModelForCausalLM
+from transformers import AutoProcessor
+from transformers import BitsAndBytesConfig
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import torch
+import pandas as pd
+from torchmetrics.text import CharErrorRate
+from peft import PeftModel, PeftConfig
+
+# Define model ID
+model_id = "microsoft/Phi-3-vision-128k-instruct"
+peft_model_id = "hadrakey/alphapen_phi3"
+peft_model_id_new = "hadrakey/alphapen_new_large"
+
+# Load processor
+processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
+
+# phi3 finetuned 
+# config = PeftConfig.from_pretrained(peft_model_id)
+
+# processor_fine = AutoProcessor.from_pretrained(config.base_model_name_or_path, trust_remote_code=True)
+
+# Finetuned model
+# config_new = PeftConfig.from_pretrained(peft_model_id_new)
+model_finetune = VisionEncoderDecoderModel.from_pretrained("hadrakey/alphapen_large")
+# model_new_finetune = AutoModelForCausalLM.from_pretrained(config_new.base_model_name_or_path, device_map="auto", trust_remote_code=True, torch_dtype="auto")
+
+# model_finetune_phi3 = AutoModelForCausalLM.from_pretrained("hadrakey/alphapen_phi3", trust_remote_code=True)
+
+#Baseline
+model_base = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
+
+
+processor_ocr = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
+# processor_ocr_new = AutoProcessor.from_pretrained(config_new.base_model_name_or_path, device_map="auto", trust_remote_code=True, torch_dtype="auto")
+# Define BitsAndBytes configuration for 4-bit quantization
+nf4_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+)
+
+# Load model with 4-bit quantization and map to CUDA
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="cuda",
+    trust_remote_code=True,
+    torch_dtype="auto",
+    quantization_config=nf4_config,
+)
+
+# base_model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, device_map="auto", trust_remote_code=True, torch_dtype="auto")
+
+# model_finetune_phi3 = PeftModel.from_pretrained(base_model, peft_model_id)
+# Define initial chat message with image placeholder
+messages = [{"role": "user", "content": """<|image_1|>\nThis image contains handwritten French characters forming a complete or partial word. The image is blurred, which makes recognition challenging. Please analyze the image to the best of your ability and provide your best guess of the French word or partial word shown, even if you're not certain. Follow these guidelines:
+
+1. Examine the overall shape and any discernible character features.
+2. Consider common French letter combinations and word patterns.
+3. If you can only identify some characters, provide those as a partial word.
+4. Make an educated guess based on what you can see, even if it's just a few letters.
+5. If you can see any characters at all, avoid responding with "indiscernible."
+
+Your response should be only the predicted French word or partial word, using lowercase letters unless capital letters are clearly visible. If you can see any characters or shapes at all, provide the OCR from the image.
+"""}]
+
+# messages = [{"role": "user", "content": """<|image_1|>\nWhat is shown is this images ? You should only output only your guess otherwise output the OCR.
+# """}]
+
+# Download image from URL
+url = "https://images.unsplash.com/photo-1528834342297-fdefb9a5a92b?ixlib=rb-4.0.3&q=85&fm=jpg&crop=entropy&cs=srgb&dl=roonz-nl-vjDbHCjHlEY-unsplash.jpg&w=640"
+# image = Image.open(requests.get(url, stream=True).raw)
+
+df_path = "/mnt/data1/Datasets/AlphaPen/" + "testing_data.csv"
+data = pd.read_csv(df_path)
+data.dropna(inplace=True)
+data.reset_index(inplace=True)
+sample = data.iloc[:5000,:]
+root_dir = "/mnt/data1/Datasets/OCR/Alphapen/clean_data/"
+# Prepare prompt with image token
+prompt = processor.tokenizer.apply_chat_template(
+    messages, tokenize=False, add_generation_prompt=True
+)
+cer_metric = CharErrorRate()
+phi_output=[]
+phi_finetune_output=[]
+inf_baseline = []
+inf_finetune = []
+inf_finetune_new = []
+
+cer_phi = []
+cer_phi_finetune = []
+cer_trocr_fine_new = []
+cer_trocr_fine = []
+cer_trocr_base = []
+for idx in range(len(sample)):
+
+# idx=30  # choose the image
+    image = Image.open(root_dir + "final_cropped_rotated_" + data.filename[idx]).convert("RGB")
+
+    # Process prompt and image for model input
+    inputs = processor(prompt, [image], return_tensors="pt").to("cuda:0")
+
+    # Generate text response using model
+    generate_ids = model.generate(
+        **inputs,
+        eos_token_id=processor.tokenizer.eos_token_id,
+        max_new_tokens=500,
+        do_sample=False,
+    )
+
+    # Remove input tokens from generated response
+    generate_ids = generate_ids[:, inputs["input_ids"].shape[1] :]
+
+    # Decode generated IDs to text
+    response = processor.batch_decode(
+        generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    )[0]
+    phi_output.append(response)
+    cer_phi.append(cer_metric(response.lower(), data.text[idx].lower()).detach().numpy())
+
+    # Generate text response using model finetuned
+    # generate_ids_fine = model_finetune_phi3.generate(
+    #     **inputs,
+    #     eos_token_id=processor.tokenizer.eos_token_id,
+    #     max_new_tokens=500,
+    #     do_sample=False,
+    # )
+
+    # # Remove input tokens from generated response
+    # inputs = processor_fine(prompt, [image], return_tensors="pt").to("cuda:0")
+    # generate_ids_fine = generate_ids_fine[:, inputs["input_ids"].shape[1] :]
+
+    # Decode generated IDs to text
+    # response = processor.batch_decode(
+    #     generate_ids_fine, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    # )[0]
+    # phi_finetune_output.append(response)
+    # cer_phi_finetune.append(cer_metric(response, data.text[idx]).detach().numpy())
+
+    # Trocr
+    pixel_values = processor_ocr(image, return_tensors="pt").pixel_values
+    generated_ids_base = model_base.generate(pixel_values)
+    generated_ids_fine = model_finetune.generate(pixel_values)
+    # generated_ids_fine_new = model_finetune_new.generate(pixel_values)
+    generated_text_base = processor_ocr.batch_decode(generated_ids_base, skip_special_tokens=True)[0]
+    generated_text_fine= processor_ocr.batch_decode(generated_ids_fine, skip_special_tokens=True)[0]
+    # generated_text_fine_new= processor_ocr_new.batch_decode(generated_ids_fine_new, skip_special_tokens=True)[0]
+    
+    inf_baseline.append(generated_text_base)
+    inf_finetune.append(generated_text_fine)
+    # inf_finetune_new.append(generated_text_fine_new)
+
+    # cer_trocr_fine_new.append(cer_metric(generated_text_fine_new, data.text[idx]).detach().numpy())
+    cer_trocr_fine.append(cer_metric(generated_text_fine.lower(), data.text[idx].lower()).detach().numpy())
+    cer_trocr_base.append(cer_metric(generated_text_base.lower(), data.text[idx].lower()).detach().numpy())
+
+
+# Print the generated response
+sample["phi3"]=phi_output
+# sample["phi3_fine"]=phi_finetune_output
+sample["Baseline"]=inf_baseline
+sample["Finetune"]=inf_finetune  
+# sample["Finetune_new"]=inf_finetune_new  
+sample["cer_phi"]=cer_phi
+# sample["cer_phi_fine"]=cer_phi_finetune
+sample["cer_trocr_base"]=cer_trocr_base
+sample["cer_trocr_fine"]=cer_trocr_fine  
+# sample["cer_trocr_fine_new"]=cer_trocr_fine_new  
+sample.to_csv("/mnt/data1/Datasets/AlphaPen/" + "sample_data.csv")