LSTM_Model

from datasets import load_dataset from transformers import AutoTokenizer import torch from torch.utils.data import DataLoader

Load English-Estonian dataset from HuggingFace

dataset = load_dataset("opus100", "en-et")

Initialize tokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") # Adjust if needed for your language model

Tokenization function with batched input handling

def tokenize_data(batch): # Extract English and Estonian translations en_texts = [example["en"] for example in batch["translation"]] et_texts = [example["et"] for example in batch["translation"]]

# Tokenize the extracted texts
inputs = tokenizer(en_texts, truncation=True, padding="max_length", max_length=128, return_tensors="pt")
targets = tokenizer(et_texts, truncation=True, padding="max_length", max_length=128, return_tensors="pt")

return {"input_ids": inputs["input_ids"].squeeze(), "labels": targets["input_ids"].squeeze()}

Tokenize datasets with batched processing

train_data = dataset["train"].map(tokenize_data, batched=True) val_data = dataset["validation"].map(tokenize_data, batched=True) test_data = dataset["test"].map(tokenize_data, batched=True)

DataLoader creation

def create_dataloader(data, batch_size=32): return DataLoader(data.with_format("torch"), batch_size=batch_size, shuffle=True)

train_dataloader = create_dataloader(train_data, batch_size=32) val_dataloader = create_dataloader(val_data, batch_size=32) test_dataloader = create_dataloader(test_data, batch_size=32)

Extract a single sample from the training data

sample = train_data[0] # Get the first example from train_data

Decode the input and target

input_text = tokenizer.decode(sample["input_ids"], skip_special_tokens=True) label_text = tokenizer.decode(sample["labels"], skip_special_tokens=True)

Print the formatted sample

print("EN to ET\n") print("<s>") print(f"<en> {input_text} </en>") print(f"<et> {label_text} </et>") print("</s>\n")

print("ET to EN\n") print("<s>") print(f"<et> {label_text} </et>") print(f"<en> {input_text} </en>") print("</s>")

import torch.nn as nn

class LSTMTranslator(nn.Module): def init(self, input_dim, hidden_dim, output_dim, n_layers=2): # Default 2 layers super(LSTMTranslator, self).init() self.embedding = nn.Embedding(input_dim, hidden_dim) self.lstm = nn.LSTM(hidden_dim, hidden_dim, num_layers=n_layers, batch_first=True) self.fc = nn.Linear(hidden_dim, output_dim)

def forward(self, src):
    embedded = self.embedding(src)
    outputs, _ = self.lstm(embedded)
    predictions = self.fc(outputs)
    return predictions

import time import math import torch import pandas as pd import evaluate from torch import nn, optim from torch.utils.data import DataLoader

Load BLEU and ChrF evaluation metrics

bleu_metric = evaluate.load("bleu") chrf_metric = evaluate.load("chrf")

class LSTMModel(nn.Module): def init(self, input_dim, embed_dim, hidden_dim, output_dim): super(LSTMModel, self).init() self.embedding = nn.Embedding(input_dim, embed_dim) self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True) self.fc = nn.Linear(hidden_dim, output_dim)

def forward(self, x):
    embedded = self.embedding(x)
    lstm_out, _ = self.lstm(embedded)
    output = self.fc(lstm_out)
    return output

Training and evaluation function

def train_lstm(model, train_dataloader, val_dataloader, criterion, optimizer, test_dataloader, tokenizer, device, n_epochs=5): model.to(device) train_losses, val_losses = [], [] bleu_scores, chrf_scores = [], []

for epoch in range(n_epochs):
    model.train()
    train_loss = 0
    start_time = time.time()

    for batch in train_dataloader:
        src = batch["input_ids"].to(device)
        tgt = batch["labels"].to(device)

        optimizer.zero_grad()

        # Forward pass
        outputs = model(src)
        tgt = tgt[:, :outputs.size(1)]  # Align target length with output length
        outputs = outputs.reshape(-1, outputs.size(-1))
        tgt = tgt.reshape(-1)

        # Compute loss
        loss = criterion(outputs, tgt)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()

    train_loss /= len(train_dataloader)
    train_losses.append(train_loss)

    # Validation phase
    model.eval()
    val_loss = 0
    with torch.no_grad():
        for batch in val_dataloader:
            src = batch["input_ids"].to(device)
            tgt = batch["labels"].to(device)

            outputs = model(src)
            tgt = tgt[:, :outputs.size(1)]
            outputs = outputs.reshape(-1, outputs.size(-1))
            tgt = tgt.reshape(-1)

            loss = criterion(outputs, tgt)
            val_loss += loss.item()

    val_loss /= len(val_dataloader)
    val_losses.append(val_loss)

    # Compute BLEU and ChrF scores on test set
    bleu_score, chrf_score = compute_bleu_chrf_scores(model, test_dataloader, tokenizer, device)
    bleu_scores.append({"Epoch": epoch + 1, "BLEU Score": bleu_score})
    chrf_scores.append({"Epoch": epoch + 1, "ChrF Score": chrf_score})

    epoch_time = time.time() - start_time
    print(f"Epoch {epoch+1}/{n_epochs} completed in {epoch_time:.2f}s")
    print(f"Train Loss: {train_loss:.4f}, Validation Loss: {val_loss:.4f}")
    print(f"BLEU Score: {bleu_score:.4f}, ChrF Score: {chrf_score:.4f}")

# Save BLEU and ChrF scores to CSV
bleu_df = pd.DataFrame(bleu_scores)
bleu_df.to_csv("lstm_bleu_scores.csv", index=False)
chrf_df = pd.DataFrame(chrf_scores)
chrf_df.to_csv("lstm_chrf_scores.csv", index=False)

print("Training completed. BLEU and ChrF scores saved to CSV files.")
return train_losses, val_losses

Compute BLEU and ChrF scores

def compute_bleu_chrf_scores(model, dataloader, tokenizer, device): model.eval() predictions, references = [], []

with torch.no_grad():
    for batch in dataloader:
        src = batch["input_ids"].to(device)
        tgt = batch["labels"].to(device)

        outputs = model(src)
        predicted_ids = outputs.argmax(dim=-1)

        predicted_texts = tokenizer.batch_decode(predicted_ids, skip_special_tokens=True)
        target_texts = tokenizer.batch_decode(tgt, skip_special_tokens=True)

        # Extend predictions and references only if valid
        if predicted_texts and target_texts:
            predictions.extend(predicted_texts)
            references.extend([[text] for text in target_texts])

# Handle cases where predictions or references are empty
if not predictions or not references:
    print("Warning: Empty predictions or references. Returning BLEU and ChrF scores as 0.")
    return 0.0, 0.0

bleu_score = bleu_metric.compute(predictions=predictions, references=references)['bleu']
chrf_score = chrf_metric.compute(predictions=predictions, references=references)['score']
return bleu_score, chrf_score

Load dataset

from datasets import load_dataset from transformers import AutoTokenizer

dataset = load_dataset("opus100", "en-et") tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

Tokenization

def tokenize_data(batch): en_texts = [example["en"] for example in batch["translation"]] et_texts = [example["et"] for example in batch["translation"]] inputs = tokenizer(en_texts, truncation=True, padding="max_length", max_length=128, return_tensors="pt") targets = tokenizer(et_texts, truncation=True, padding="max_length", max_length=128, return_tensors="pt") return {"input_ids": inputs["input_ids"], "labels": targets["input_ids"]}

train_data = dataset["train"].map(tokenize_data, batched=True) val_data = dataset["validation"].map(tokenize_data, batched=True) test_data = dataset["test"].map(tokenize_data, batched=True)

def create_dataloader(data, batch_size=32): return DataLoader(data.with_format("torch"), batch_size=batch_size, shuffle=True)

train_dataloader = create_dataloader(train_data) val_dataloader = create_dataloader(val_data) test_dataloader = create_dataloader(test_data)

Model, optimizer, and criterion

input_dim = tokenizer.vocab_size embed_dim = 128 hidden_dim = 256 output_dim = tokenizer.vocab_size

lstm_model = LSTMModel(input_dim, embed_dim, hidden_dim, output_dim) criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id) optimizer = optim.Adam(lstm_model.parameters(), lr=0.001)

Train model

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") train_losses, val_losses = train_lstm( lstm_model, train_dataloader, val_dataloader, criterion, optimizer, test_dataloader, tokenizer, device, n_epochs=5 )

Save Model State

torch.save(lstm_model.state_dict(), "lstm_model_state_dict.pth")

Define the exact same LSTMModel class for loading

class LSTMModel(nn.Module): def init(self, input_dim, embed_dim, hidden_dim, output_dim): super(LSTMModel, self).init() self.embedding = nn.Embedding(input_dim, embed_dim) self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True) self.fc = nn.Linear(hidden_dim, output_dim)

def forward(self, x):
    embedded = self.embedding(x)
    lstm_out, _ = self.lstm(embedded)
    output = self.fc(lstm_out)
    return output

Reinitialize model with the exact same parameters

loaded_model = LSTMModel(input_dim, embed_dim, hidden_dim, output_dim)

Load saved state dict

loaded_model.load_state_dict(torch.load("lstm_model_state_dict.pth"))

Set to evaluation mode for inference

loaded_model.eval()

print("Model successfully loaded and set to evaluation mode.")