diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..cd703ca
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,4 @@
+env
+Donut
+nohup.out
+wandb
diff --git a/donut-eval.py b/donut-eval.py
new file mode 100644
index 0000000..c5548e6
--- /dev/null
+++ b/donut-eval.py
@@ -0,0 +1,83 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[1]:
+
+
+from transformers import DonutProcessor, VisionEncoderDecoderModel
+from datasets import load_dataset
+import re
+import json
+import torch
+from tqdm.auto import tqdm
+import numpy as np
+
+from donut import JSONParseEvaluator
+
+
+# In[2]:
+
+
+processor = DonutProcessor.from_pretrained("Zombely/plwiki-proto-fine-tuned")
+model = VisionEncoderDecoderModel.from_pretrained("Zombely/plwiki-proto-fine-tuned")
+
+
+# In[3]:
+
+
+dataset = load_dataset("Zombely/pl-text-images-5000-whole", split="validation")
+
+
+# In[4]:
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+model.eval()
+model.to(device)
+
+output_list = []
+accs = []
+
+
+for idx, sample in tqdm(enumerate(dataset), total=len(dataset)):
+    # prepare encoder inputs
+    pixel_values = processor(sample["image"].convert("RGB"), return_tensors="pt").pixel_values
+    pixel_values = pixel_values.to(device)
+    # prepare decoder inputs
+    task_prompt = "<s_cord-v2>"
+    decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids
+    decoder_input_ids = decoder_input_ids.to(device)
+    
+    # autoregressively generate sequence
+    outputs = model.generate(
+            pixel_values,
+            decoder_input_ids=decoder_input_ids,
+            max_length=model.decoder.config.max_position_embeddings,
+            early_stopping=True,
+            pad_token_id=processor.tokenizer.pad_token_id,
+            eos_token_id=processor.tokenizer.eos_token_id,
+            use_cache=True,
+            num_beams=1,
+            bad_words_ids=[[processor.tokenizer.unk_token_id]],
+            return_dict_in_generate=True,
+        )
+
+    # turn into JSON
+    seq = processor.batch_decode(outputs.sequences)[0]
+    seq = seq.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "")
+    seq = re.sub(r"<.*?>", "", seq, count=1).strip()  # remove first task start token
+    seq = processor.token2json(seq)
+
+    ground_truth = json.loads(sample["ground_truth"])
+    ground_truth = ground_truth["gt_parse"]
+    evaluator = JSONParseEvaluator()
+    score = evaluator.cal_acc(seq, ground_truth)
+
+    accs.append(score)
+    output_list.append(seq)
+
+scores = {"accuracies": accs, "mean_accuracy": np.mean(accs)}
+print(scores, f"length : {len(accs)}")
+print("Mean accuracy:", np.mean(accs))
+
diff --git a/donut-train.py b/donut-train.py
new file mode 100644
index 0000000..6590e1d
--- /dev/null
+++ b/donut-train.py
@@ -0,0 +1,387 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[19]:
+
+
+from transformers import VisionEncoderDecoderConfig, DonutProcessor, VisionEncoderDecoderModel
+from datasets import load_dataset
+import json
+import random
+from typing import Any, List, Tuple
+import torch
+from torch.utils.data import Dataset, DataLoader
+import re
+from nltk import edit_distance
+import numpy as np
+from pytorch_lightning.loggers import WandbLogger
+from pytorch_lightning.callbacks import Callback
+import pytorch_lightning as pl
+import os
+from huggingface_hub import login
+
+
+# In[8]:
+
+
+DATASET_PATH = "Zombely/pl-text-images-5000-whole"
+PRETRAINED_MODEL_PATH = "Zombely/plwiki-proto-fine-tuned"
+START_MODEL_PATH = "Zombely/plwiki-proto-fine-tuned"
+OUTPUT_MODEL_PATH = "Zombely/plwiki-proto-fine-tuned-v2"
+LOGGING_PATH = "plwiki-proto-ft-second-iter"
+
+
+# In[ ]:
+
+
+train_config = {
+    "max_epochs":30,
+    "val_check_interval":0.5, # how many times we want to validate during an epoch
+    "check_val_every_n_epoch":1,
+    "gradient_clip_val":1.0,
+    "num_training_samples_per_epoch": 800,
+    "lr":3e-5,
+    "train_batch_sizes": [8],
+    "val_batch_sizes": [1],
+    # "seed":2022,
+    "num_nodes": 1,
+    "warmup_steps": 300, # 800/8*30/10, 10%
+    "result_path": "./result",
+    "verbose": True,
+}
+
+
+# In[9]:
+
+
+dataset = load_dataset(DATASET_PATH)
+
+
+# In[10]:
+
+
+max_length = 768
+image_size = [1920, 2560]
+config = VisionEncoderDecoderConfig.from_pretrained(PRETRAINED_MODEL_PATH)
+config.encoder.image_size = image_size # (height, width)
+config.decoder.max_length = max_length
+
+
+# In[11]:
+
+
+processor = DonutProcessor.from_pretrained(START_MODEL_PATH)
+model = VisionEncoderDecoderModel.from_pretrained(PRETRAINED_MODEL_PATH, config=config)
+
+
+# In[12]:
+
+
+added_tokens = []
+
+class DonutDataset(Dataset):
+    """
+    DonutDataset which is saved in huggingface datasets format. (see details in https://huggingface.co/docs/datasets)
+    Each row, consists of image path(png/jpg/jpeg) and gt data (json/jsonl/txt),
+    and it will be converted into input_tensor(vectorized image) and input_ids(tokenized string).
+    Args:
+        dataset_name_or_path: name of dataset (available at huggingface.co/datasets) or the path containing image files and metadata.jsonl
+        max_length: the max number of tokens for the target sequences
+        split: whether to load "train", "validation" or "test" split
+        ignore_id: ignore_index for torch.nn.CrossEntropyLoss
+        task_start_token: the special token to be fed to the decoder to conduct the target task
+        prompt_end_token: the special token at the end of the sequences
+        sort_json_key: whether or not to sort the JSON keys
+    """
+
+    def __init__(
+        self,
+        dataset_name_or_path: str,
+        max_length: int,
+        split: str = "train",
+        ignore_id: int = -100,
+        task_start_token: str = "<s>",
+        prompt_end_token: str = None,
+        sort_json_key: bool = True,
+    ):
+        super().__init__()
+
+        self.max_length = max_length
+        self.split = split
+        self.ignore_id = ignore_id
+        self.task_start_token = task_start_token
+        self.prompt_end_token = prompt_end_token if prompt_end_token else task_start_token
+        self.sort_json_key = sort_json_key
+
+        self.dataset = load_dataset(dataset_name_or_path, split=self.split)
+        self.dataset_length = len(self.dataset)
+
+        self.gt_token_sequences = []
+        for sample in self.dataset:
+            ground_truth = json.loads(sample["ground_truth"])
+            if "gt_parses" in ground_truth:  # when multiple ground truths are available, e.g., docvqa
+                assert isinstance(ground_truth["gt_parses"], list)
+                gt_jsons = ground_truth["gt_parses"]
+            else:
+                assert "gt_parse" in ground_truth and isinstance(ground_truth["gt_parse"], dict)
+                gt_jsons = [ground_truth["gt_parse"]]
+
+            self.gt_token_sequences.append(
+                [
+                    self.json2token(
+                        gt_json,
+                        update_special_tokens_for_json_key=self.split == "train",
+                        sort_json_key=self.sort_json_key,
+                    )
+                    + processor.tokenizer.eos_token
+                    for gt_json in gt_jsons  # load json from list of json
+                ]
+            )
+
+        self.add_tokens([self.task_start_token, self.prompt_end_token])
+        self.prompt_end_token_id = processor.tokenizer.convert_tokens_to_ids(self.prompt_end_token)
+
+    def json2token(self, obj: Any, update_special_tokens_for_json_key: bool = True, sort_json_key: bool = True):
+        """
+        Convert an ordered JSON object into a token sequence
+        """
+        if type(obj) == dict:
+            if len(obj) == 1 and "text_sequence" in obj:
+                return obj["text_sequence"]
+            else:
+                output = ""
+                if sort_json_key:
+                    keys = sorted(obj.keys(), reverse=True)
+                else:
+                    keys = obj.keys()
+                for k in keys:
+                    if update_special_tokens_for_json_key:
+                        self.add_tokens([fr"<s_{k}>", fr"</s_{k}>"])
+                    output += (
+                        fr"<s_{k}>"
+                        + self.json2token(obj[k], update_special_tokens_for_json_key, sort_json_key)
+                        + fr"</s_{k}>"
+                    )
+                return output
+        elif type(obj) == list:
+            return r"<sep/>".join(
+                [self.json2token(item, update_special_tokens_for_json_key, sort_json_key) for item in obj]
+            )
+        else:
+            obj = str(obj)
+            if f"<{obj}/>" in added_tokens:
+                obj = f"<{obj}/>"  # for categorical special tokens
+            return obj
+    
+    def add_tokens(self, list_of_tokens: List[str]):
+        """
+        Add special tokens to tokenizer and resize the token embeddings of the decoder
+        """
+        newly_added_num = processor.tokenizer.add_tokens(list_of_tokens)
+        if newly_added_num > 0:
+            model.decoder.resize_token_embeddings(len(processor.tokenizer))
+            added_tokens.extend(list_of_tokens)
+    
+    def __len__(self) -> int:
+        return self.dataset_length
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Load image from image_path of given dataset_path and convert into input_tensor and labels
+        Convert gt data into input_ids (tokenized string)
+        Returns:
+            input_tensor : preprocessed image
+            input_ids : tokenized gt_data
+            labels : masked labels (model doesn't need to predict prompt and pad token)
+        """
+        sample = self.dataset[idx]
+
+        # inputs
+        pixel_values = processor(sample["image"], random_padding=self.split == "train", return_tensors="pt").pixel_values
+        pixel_values = pixel_values.squeeze()
+
+        # targets
+        target_sequence = random.choice(self.gt_token_sequences[idx])  # can be more than one, e.g., DocVQA Task 1
+        input_ids = processor.tokenizer(
+            target_sequence,
+            add_special_tokens=False,
+            max_length=self.max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt",
+        )["input_ids"].squeeze(0)
+
+        labels = input_ids.clone()
+        labels[labels == processor.tokenizer.pad_token_id] = self.ignore_id  # model doesn't need to predict pad token
+        # labels[: torch.nonzero(labels == self.prompt_end_token_id).sum() + 1] = self.ignore_id  # model doesn't need to predict prompt (for VQA)
+        return pixel_values, labels, target_sequence
+
+
+# In[13]:
+
+
+processor.image_processor.size = image_size[::-1] # should be (width, height)
+processor.image_processor.do_align_long_axis = False
+
+train_dataset = DonutDataset(DATASET_PATH, max_length=max_length,
+                             split="train", task_start_token="<s_cord-v2>", prompt_end_token="<s_cord-v2>",
+                             sort_json_key=False, # cord dataset is preprocessed, so no need for this
+                             )
+
+val_dataset = DonutDataset(DATASET_PATH, max_length=max_length,
+                             split="validation", task_start_token="<s_cord-v2>", prompt_end_token="<s_cord-v2>",
+                             sort_json_key=False, # cord dataset is preprocessed, so no need for this
+                             )
+
+
+# In[14]:
+
+
+model.config.pad_token_id = processor.tokenizer.pad_token_id
+model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s_cord-v2>'])[0]
+
+
+# In[15]:
+
+
+train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=4)
+val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=4)
+
+
+# In[16]:
+
+
+class DonutModelPLModule(pl.LightningModule):
+    def __init__(self, config, processor, model):
+        super().__init__()
+        self.config = config
+        self.processor = processor
+        self.model = model
+
+    def training_step(self, batch, batch_idx):
+        pixel_values, labels, _ = batch
+        
+        outputs = self.model(pixel_values, labels=labels)
+        loss = outputs.loss
+        self.log_dict({"train_loss": loss}, sync_dist=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx, dataset_idx=0):
+        pixel_values, labels, answers = batch
+        batch_size = pixel_values.shape[0]
+        # we feed the prompt to the model
+        decoder_input_ids = torch.full((batch_size, 1), self.model.config.decoder_start_token_id, device=self.device)
+        
+        outputs = self.model.generate(pixel_values,
+                                   decoder_input_ids=decoder_input_ids,
+                                   max_length=max_length,
+                                   early_stopping=True,
+                                   pad_token_id=self.processor.tokenizer.pad_token_id,
+                                   eos_token_id=self.processor.tokenizer.eos_token_id,
+                                   use_cache=True,
+                                   num_beams=1,
+                                   bad_words_ids=[[self.processor.tokenizer.unk_token_id]],
+                                   return_dict_in_generate=True,)
+    
+        predictions = []
+        for seq in self.processor.tokenizer.batch_decode(outputs.sequences):
+            seq = seq.replace(self.processor.tokenizer.eos_token, "").replace(self.processor.tokenizer.pad_token, "")
+            seq = re.sub(r"<.*?>", "", seq, count=1).strip()  # remove first task start token
+            predictions.append(seq)
+
+        scores = list()
+        for pred, answer in zip(predictions, answers):
+            pred = re.sub(r"(?:(?<=>) | (?=</s_))", "", pred)
+            # NOT NEEDED ANYMORE
+            # answer = re.sub(r"<.*?>", "", answer, count=1)
+            answer = answer.replace(self.processor.tokenizer.eos_token, "")
+            scores.append(edit_distance(pred, answer) / max(len(pred), len(answer)))
+
+            if self.config.get("verbose", False) and len(scores) == 1:
+                print(f"Prediction: {pred}")
+                print(f"    Answer: {answer}")
+                print(f" Normed ED: {scores[0]}")
+
+        return scores
+
+    def validation_epoch_end(self, validation_step_outputs):
+        # I set this to 1 manually
+        # (previously set to len(self.config.dataset_name_or_paths))
+        num_of_loaders = 1
+        if num_of_loaders == 1:
+            validation_step_outputs = [validation_step_outputs]
+        assert len(validation_step_outputs) == num_of_loaders
+        cnt = [0] * num_of_loaders
+        total_metric = [0] * num_of_loaders
+        val_metric = [0] * num_of_loaders
+        for i, results in enumerate(validation_step_outputs):
+            for scores in results:
+                cnt[i] += len(scores)
+                total_metric[i] += np.sum(scores)
+            val_metric[i] = total_metric[i] / cnt[i]
+            val_metric_name = f"val_metric_{i}th_dataset"
+            self.log_dict({val_metric_name: val_metric[i]}, sync_dist=True)
+        self.log_dict({"val_metric": np.sum(total_metric) / np.sum(cnt)}, sync_dist=True)
+
+    def configure_optimizers(self):
+        # TODO add scheduler
+        optimizer = torch.optim.Adam(self.parameters(), lr=self.config.get("lr"))
+    
+        return optimizer
+
+    def train_dataloader(self):
+        return train_dataloader
+
+    def val_dataloader(self):
+        return val_dataloader
+
+
+# In[17]:
+
+
+class PushToHubCallback(Callback):
+    def on_train_epoch_end(self, trainer, pl_module):
+        print(f"Pushing model to the hub, epoch {trainer.current_epoch}")
+        pl_module.model.push_to_hub(OUTPUT_MODEL_PATH,
+                                    commit_message=f"Training in progress, epoch {trainer.current_epoch}")
+        pl_module.processor.push_to_hub(OUTPUT_MODEL_PATH, commit_message=f"Training in progress, epoch {trainer.current_epoch}")
+
+    def on_train_end(self, trainer, pl_module):
+        print(f"Pushing model to the hub after training")
+        pl_module.processor.push_to_hub(OUTPUT_MODEL_PATH,
+                                    commit_message=f"Training done")
+        pl_module.model.push_to_hub(OUTPUT_MODEL_PATH,
+                                    commit_message=f"Training done")
+
+
+# In[18]:
+
+
+login(os.environ.get("HUG_TOKKEN", ""))
+
+
+# ### Wandb.ai link: https://wandb.ai/michalkozlowski936/Donut?workspace=user-michalkozlowski936
+# ### Hugging_face link https://huggingface.co/Zombely
+
+# In[22]:
+
+
+model_module = DonutModelPLModule(train_config, processor, model)
+
+wandb_logger = WandbLogger(project="Donut", name=LOGGING_PATH)
+
+trainer = pl.Trainer(
+        accelerator="gpu", # change to gpu
+        devices=1,
+        max_epochs=train_config.get("max_epochs"),
+        val_check_interval=train_config.get("val_check_interval"),
+        check_val_every_n_epoch=train_config.get("check_val_every_n_epoch"),
+        gradient_clip_val=train_config.get("gradient_clip_val"),
+        precision=16, # we'll use mixed precision
+        num_sanity_val_steps=0,
+        logger=wandb_logger,
+        callbacks=[PushToHubCallback()],
+)
+
+trainer.fit(model_module)
+
diff --git a/requirements.txt b/requirements.txt
index 8d66d90..f1de7a0 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -14,7 +14,7 @@ beautifulsoup4==4.11.1
 bleach==5.0.1
 blend-modes==2.1.0
 cachetools==5.2.0
-certifi @ file:///croot/certifi_1665076670883/work/certifi
+certifi
 cffi==1.15.1
 charset-normalizer==2.1.1
 click==8.1.3
@@ -26,7 +26,7 @@ decorator==5.1.1
 defusedxml==0.7.1
 dill==0.3.6
 docker-pycreds==0.4.0
-donut-python @ file:///home/pc/work/donut
+donut-python
 entrypoints==0.4
 evaluate==0.3.0
 fastapi==0.87.0