from tqdm import tqdm
import Levenshtein
import torch.optim as optim
# Initialize model, loss function, and optimizer
num_classes = len(charset) + 1  # +1 for the blank label required for CTC
device = torch.device('cuda')
model = CRNN(num_classes).to(device)
criterion = nn.CTCLoss(blank=0, zero_infinity=True)
optimizer = optim.Adam(model.parameters(), lr=0.0001)

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)

def train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs=10):
    for epoch in range(num_epochs):
        model.train()
        total_loss = 0
        progress_bar = tqdm(train_loader, desc=f"Epoch [{epoch+1}/{num_epochs}]", leave=False)
        for images, labels, label_lengths in progress_bar:
            images = images.to(device)
            labels = labels.to(device)
            label_lengths = label_lengths.to(device)
            optimizer.zero_grad()
            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)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()

            progress_bar.set_postfix(loss=loss.item())

        avg_loss = total_loss / len(train_loader)
        print(f"Epoch [{epoch+1}/{num_epochs}], Training Loss: {avg_loss:.4f}")

        # Validation
        model.eval()
        with torch.no_grad():
            cer_total = 0
            wer_total = 0
            num_samples = 0
            for images, labels, label_lengths in test_loader:
                images = images.to(device)
                labels = labels.to(device)
                label_lengths = label_lengths.to(device)
                outputs = model(images)

                # Decode predictions
                decoded_preds = decode_predictions(outputs)
                # Prepare ground truth texts
                ground_truths = []
                index = 0
                labels = labels.cpu()
                for length in label_lengths.cpu():
                    gt_text = ''.join([idx_to_char.get(idx.item(), '') for idx in labels[index:index+length]])
                    ground_truths.append(gt_text)
                    index += length

                # Calculate CER and WER
                for pred, gt in zip(decoded_preds, ground_truths):
                    cer_score = cer(pred, gt)
                    wer_score = wer(pred, gt)
                    cer_total += cer_score
                    wer_total += wer_score
                    num_samples += 1

            avg_cer = cer_total / num_samples
            avg_wer = wer_total / num_samples
            print(f"Validation CER: {avg_cer:.4f}, WER: {avg_wer:.4f}")

    print("Training completed.")

train_model(model, train_loader, test_loader, criterion, optimizer, num_epochs=10)
torch.save(model.state_dict(), 'crnn_model.pth')

how to make it faster using all the cores hp victus rtx 3050