# -*- coding: utf-8 -*-

#!pip install transformers accelerate

!pip install --upgrade  transformers==4.35.2 accelerate==0.24.1
!pip install peft==0.4.0

#!pip install -q datasets jiwer
!pip install -q datasets==2.21.0 jiwer

from google.colab import drive
drive.mount('/content/drive')



import pandas as pd

column_names = ['file_name', 'text']
df = pd.read_csv('/content/drive/MyDrive/mydata_training/mydata8/data.csv',index_col=False,header=0)
del df['id']
del df['image_name']
df['word'] = df['word'].apply(lambda x: x[2:-2])
print(df.head())
print(len(df))
#df.drop(df.tail(200).index,inplace = True)
#print(len(df))

from sklearn.model_selection import train_test_split

train_df, test_df = train_test_split(df, test_size=0.1)
# we reset the indices to start from zero
train_df.reset_index(drop=True, inplace=True)
test_df.reset_index(drop=True, inplace=True)

import torch
from torchvision.transforms import Resize, ToTensor
from torch.utils.data import Dataset
from PIL import Image

class IAMDataset(Dataset):
    def __init__(self, root_dir, df, processor, max_target_length=128):
        self.root_dir = root_dir
        self.df = df
        self.processor = processor
        self.max_target_length = max_target_length

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        # get file name + text
        file_name = self.df['filename'][idx]
        text = self.df['word'][idx].rstrip()
        # prepare image (i.e. resize + normalize)
        image = Image.open(self.root_dir + file_name).convert("RGB")
        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)}
        #encoding = {"pixel_values": pixel_values, "labels": torch.tensor(labels)}
        return encoding



from transformers import TrOCRProcessor

processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten") #base-stage1

train_dataset = IAMDataset(root_dir='/content/drive/MyDrive/mydata_training/mydata8/images/',
                           df=train_df,
                           processor=processor)
eval_dataset = IAMDataset(root_dir='/content/drive/MyDrive/mydata_training/mydata8/images/',
                           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)

print("Number of training examples:", len(train_dataset))
print("Number of validation examples:", len(eval_dataset))

image = Image.open(train_dataset.root_dir + train_df['filename'][0]).convert("RGB")
image

labels = encoding['labels']
labels[labels == -100] = processor.tokenizer.pad_token_id
label_str = processor.decode(labels, skip_special_tokens=True)
print(label_str)

from transformers import VisionEncoderDecoderModel, AutoModel

model = VisionEncoderDecoderModel.from_pretrained("raxtemur/trocr-base-ru")
#model = VisionEncoderDecoderModel.from_pretrained("cyrillic-trocr/trocr-handwritten-cyrillic", early_stopping=True)



# 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 = 32
model.config.early_stopping = True
model.config.no_repeat_ngram_size = 3
model.config.length_penalty = 2.0
model.config.num_beams = 8



from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments

training_args = Seq2SeqTrainingArguments(
    predict_with_generate=True,
    evaluation_strategy="steps",
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    fp16=True,
    output_dir="/content/drive/MyDrive/SAVE_CHECK_POINT_my_trocr_model",
    logging_steps=2,
    save_steps=400,
    eval_steps=1200,
    num_train_epochs=3,
    report_to="none"  # Disable wandb integration
)

# from datasets import load_metric

# cer_metric = load_metric("cer")
!pip install evaluate
from evaluate import load
cer_metric = load("cer")

def compute_metrics(pred):
    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)

    # Convert to lowercase for case-insensitive comparison
    pred_str = [s.lower() for s in processor.batch_decode(pred_ids, skip_special_tokens=True)]
    labels_ids[labels_ids == -100] = processor.tokenizer.pad_token_id
    label_str = [s.lower() for s in processor.batch_decode(labels_ids, skip_special_tokens=True)]


    cer = cer_metric.compute(predictions=pred_str, references=label_str)
    # Identify and store bad predictions
    bad_predictions = [(p, l) for p, l in zip(pred_str, label_str) if p != l]
    good_predictions = [(p, l) for p, l in zip(pred_str, label_str) if p == l]
    print("Bad")
    print(bad_predictions[0:20])
    print("good")
    print(good_predictions[0:15])
    print(len(bad_predictions))

    return {"cer": cer}

from transformers import default_data_collator

# instantiate 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,
)
trainer.train()

# Your existing training code here

# After training, save to Google Drive
final_output_dir = "/content/drive/MyDrive/my_trocr_model"
trainer.save_model(final_output_dir)
processor.save_pretrained(final_output_dir)

print(f"Model and processor saved to Google Drive: {final_output_dir}")