import torch
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp
from torch.utils.data import Dataset, DataLoader, DistributedSampler
from torchvision import transforms
from PIL import Image
import requests
from io import BytesIO
import pandas as pd
from sklearn.model_selection import train_test_split
import numpy as np
import re
import Levenshtein
from symspellpy import SymSpell, Verbosity
import torch_xla
import torch_xla.core.xla_model as xm
import torch_xla.distributed.xla_multiprocessing as xmp
import torch_xla.distributed.parallel_loader as pl
import torch_xla.utils.utils as xu
charset = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ .,%'
char_to_idx = {char: idx + 1 for idx, char in enumerate(charset)}
idx_to_char = {idx: char for char, idx in char_to_idx.items()}

def _mp_fn(rank, flags):
    # Set the seed for each process
    torch.manual_seed(42 + rank)

    # Check for TPU availability
    device = xm.xla_device()

    class CRNN(nn.Module):
        def init(self, num_classes, input_channels=1):
            super(CRNN, self).init()

            self.cnn = nn.Sequential(
                nn.Conv2d(input_channels, 64, kernel_size=3, padding=1),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(2, 2),
                nn.Dropout(0.2),

                nn.Conv2d(64, 128, kernel_size=3, padding=1),
                nn.ReLU(inplace=True),
                nn.MaxPool2d(2, 2),

                nn.Conv2d(128, 256, kernel_size=3, padding=1),
                nn.BatchNorm2d(256),
                nn.ReLU(inplace=True),
                nn.MaxPool2d((2, 1)),

                nn.Conv2d(256, 512, kernel_size=3, padding=1),
                nn.BatchNorm2d(512),
                nn.ReLU(inplace=True),
                nn.MaxPool2d((2, 1)),

                nn.Conv2d(512, 512, kernel_size=2),
                nn.ReLU(inplace=True)
            )

            self.lstm = nn.GRU(
                input_size=512,
                hidden_size=256,
                num_layers=2,
                bidirectional=True,
                batch_first=True
            )

            self.fc = nn.Linear(512, num_classes)

        def forward(self, x):
            conv = self.cnn(x)
            b, c, h, w = conv.size()
            assert h == 1, f"Expected height == 1, got {h}"
            conv = conv.squeeze(2)
            conv = conv.permute(0, 2, 1)
            rnn_out, _ = self.lstm(conv)
            output = self.fc(rnn_out)
            return output.permute(1, 0, 2)
class OCRDataset(Dataset):
        def init(self, data_frame, transform=None):
            self.data_frame = data_frame.reset_index(drop=True)
            self.transform = transform

        def len(self):
            return len(self.data_frame)

        def getitem(self, idx):
            img_url = self.data_frame.iloc[idx]['image_link']
            label = self.data_frame.iloc[idx]['entity_value']

            # Attempt to load the image
            try:
                response = requests.get(img_url, timeout=5)  # Set a timeout to avoid long waits
                img = Image.open(BytesIO(response.content)).convert('L')  # Convert to grayscale
            except Exception as e:
                # Skip the failed image and return None to indicate skipping
                print(f"Failed to load image at index {idx}, skipping: {e}")
                return None

            # Apply transformations if available
            if self.transform:
                img = self.transform(img)

            # Convert the label to a tensor of character indices
            label_idx = [char_to_idx.get(char, 0) for char in label]
            label_tensor = torch.tensor(label_idx, dtype=torch.long)

            return img, label_tensor, len(label_idx)

    def collate_fn(batch):
        batch = [b for b in batch if b is not None]  # Filter out None entries
        if not batch:
            return None, None, None  # Return None if batch is empty
        images, labels, label_lengths = zip(*batch)
        images = torch.stack(images)
        labels = torch.cat(labels)
        label_lengths = torch.tensor(label_lengths, dtype=torch.long)
        return images, labels, label_lengths

    # Data augmentation and preprocessing transformations
    transform = transforms.Compose([
        transforms.Resize((32, 128)),
        transforms.RandomRotation(degrees=5),
        transforms.ColorJitter(brightness=0.5, contrast=0.5),
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,)),
        transforms.RandomErasing(p=0.1)
    ])

    # Define the CER and WER functions
    def cer(prediction, ground_truth):
        prediction = prediction.replace(' ', '')
        ground_truth = ground_truth.replace(' ', '')
        distance = Levenshtein.distance(prediction, ground_truth)
        return distance / max(len(ground_truth), 1)

    def wer(prediction, ground_truth):
        pred_words = prediction.split()
        gt_words = ground_truth.split()
        distance = Levenshtein.distance(' '.join(pred_words), ' '.join(gt_words))
        return distance / max(len(gt_words), 1)

    class EarlyStopping:
        def init(self, patience=5, min_delta=0):
            self.patience = patience
            self.min_delta = min_delta
            self.counter = 0
            self.best_loss = None
            self.early_stop = False

        def call(self, val_loss):
            if self.best_loss is None:
                self.best_loss = val_loss
            elif val_loss > self.best_loss - self.min_delta:
                self.counter += 1
                if self.counter >= self.patience:
                    self.early_stop = True
            else:
                self.best_loss = val_loss
                self.counter = 0
    # Load dataset (assuming a dataframe with 'image_link' and 'entity_value')
    df = pd.read_csv('/kaggle/input/dataset/shortened_dataset.csv')
    train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)

    # Create datasets
    train_dataset = OCRDataset(train_df, transform=transform)
    test_dataset = OCRDataset(test_df, transform=transform)

    # Create samplers
    train_sampler = DistributedSampler(
        train_dataset,
        num_replicas=xm.xrt_world_size(),
        rank=xm.get_ordinal(),
        shuffle=True)

    test_sampler = DistributedSampler(
        test_dataset,
        num_replicas=xm.xrt_world_size(),
        rank=xm.get_ordinal(),
        shuffle=False)

    # Create data loaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=32,
        sampler=train_sampler,
        collate_fn=collate_fn,
        num_workers=0,  # Set to 0 to avoid multiprocessing issues
        drop_last=True)

    test_loader = DataLoader(
        test_dataset,
        batch_size=32,
        sampler=test_sampler,
        collate_fn=collate_fn,
        num_workers=0,  # Set to 0 to avoid multiprocessing issues
        drop_last=False)

    # Initialize model, loss, and optimizer
    model = CRNN(num_classes=len(charset)+1).to(device)
    # Note: DDP is not typically needed with torch_xla on TPUs
    # criterion = nn.CTCLoss(blank=0, zero_infinity=True).to(device)
    criterion = nn.CTCLoss(blank=0).to(device)
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    def decode_output(output):
        _, preds = torch.max(output, 2)
        preds = preds.permute(1, 0)  # [batch_size, seq_len]
        pred_strings = []
        for pred in preds:
            pred_indices = pred.cpu().numpy()
            pred_indices = [p for i, p in enumerate(pred_indices) if (p != 0 and (i == 0 or p != pred_indices[i - 1]))]
            pred_str = ''.join([idx_to_char.get(idx, '') for idx in pred_indices])
            pred_strings.append(pred_str)
        return pred_strings

    def decode_label(label):
        label_indices = label.cpu().numpy()
        label_str = ''.join([idx_to_char.get(idx, '') for idx in label_indices])
        return label_str

    def train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs=20):
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2, verbose=True)
        early_stopping = EarlyStopping(patience=5)

        for epoch in range(num_epochs):
            model.train()
            total_loss = 0.0
            train_sampler.set_epoch(epoch)  # Set epoch for sampler
            # Wrap your data loader with ParallelLoader
            para_train_loader = pl.ParallelLoader(train_loader, [device]).per_device_loader(device)

            for batch_idx, batch in enumerate(para_train_loader):
                if batch[0] is None:
                    continue  # Skip empty batches
                images, labels, label_lengths = batch

                # Move data to device
                images = images.to(device)
                labels = labels.to(device)
                label_lengths = label_lengths.to(device)

                optimizer.zero_grad()
                outputs = model(images)

                # Apply log softmax to outputs before passing to CTCLoss
                outputs = torch.nn.functional.log_softmax(outputs, dim=-1)

                batch_size = images.size(0)
                seq_len = outputs.size(0)
                input_lengths = torch.full(size=(batch_size,), fill_value=seq_len, dtype=torch.long).to(device)

                loss = criterion(outputs, labels, input_lengths, label_lengths)
                if torch.isnan(loss) or torch.isinf(loss):
                    print(f"Found invalid loss at epoch {epoch+1}, skipping this batch")
                    continue  # Skip if nan or inf loss

                loss.backward()
                xm.optimizer_step(optimizer)  # Use xm.optimizer_step for TPU

                total_loss += loss.item()
                if batch_idx % 10 == 0:
                    xm.master_print(f"Epoch [{epoch+1}/{num_epochs}], Batch [{batch_idx+1}], Loss: {loss.item():.4f}")

            # Average loss across processes
            total_loss_reduced = xm.mesh_reduce('train_loss_reduce', total_loss, np.mean)
            avg_train_loss = total_loss_reduced / len(train_loader)
            xm.master_print(f"Epoch [{epoch+1}/{num_epochs}], Training Loss: {avg_train_loss:.4f}")

            # Validate the model
            val_loss, val_cer, val_wer = validate_model(model, test_loader, criterion)
            xm.master_print(f"Validation Loss: {val_loss:.4f}, CER: {val_cer:.4f}, WER: {val_wer:.4f}")

            # Scheduler step
            scheduler.step(val_loss)
            early_stopping(val_loss)
            if early_stopping.early_stop:
                xm.master_print("Early stopping triggered")
                break

    def validate_model(model, test_loader, criterion):
        model.eval()
        total_loss = 0.0
        total_cer = 0.0
        total_wer = 0.0
        count = 0
        # Wrap test_loader with ParallelLoader
        para_test_loader = pl.ParallelLoader(test_loader, [device]).per_device_loader(device)
        with torch.no_grad():
            for batch in para_test_loader:
                if batch[0] is None:
                    continue  # Skip empty batches
                images, labels, label_lengths = batch
                images = images.to(device)
                labels = labels.to(device)
                label_lengths = label_lengths.to(device)

                outputs = model(images)

                batch_size = images.size(0)
                seq_len = outputs.size(0)
                input_lengths = torch.full(size=(batch_size,), fill_value=seq_len, dtype=torch.long).to(device)

                loss = criterion(outputs, labels, input_lengths, label_lengths)
                total_loss += loss.item()

                pred_strings = decode_output(outputs)
                for i in range(batch_size):
                    pred_str = pred_strings[i]
                    gt_str = decode_label(labels[count:count+label_lengths[i]])
                    total_cer += cer(pred_str, gt_str)
                    total_wer += wer(pred_str, gt_str)
                    count += label_lengths[i]

        # Aggregate metrics across processes
        total_loss_reduced = xm.mesh_reduce('val_loss_reduce', total_loss, np.mean)
        total_cer_reduced = xm.mesh_reduce('val_cer_reduce', total_cer, np.mean)
        total_wer_reduced = xm.mesh_reduce('val_wer_reduce', total_wer, np.mean)
        count_reduced = xm.mesh_reduce('count_reduce', count, np.sum)

        avg_loss = total_loss_reduced / len(test_loader)
        avg_cer = total_cer_reduced / count_reduced
        avg_wer = total_wer_reduced / count_reduced
        return avg_loss, avg_cer, avg_wer

    # Now, start training
    train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs=20)
if name == 'main':
    FLAGS = {}
    xmp.spawn(_mp_fn, args=(FLAGS,), nprocs=8, start_method='fork')

what is wrong in the code? 
TypeError                                 Traceback (most recent call last)
Cell In[6], line 6
      3 train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)
      5 # Create datasets
----> 6 train_dataset = OCRDataset(train_df, transform=transform)
      7 test_dataset = OCRDataset(test_df, transform=transform)
      9 # Create samplers

TypeError: OCRDataset() takes no arguments