mT5 › Working with mT5

import torch
from transformers import AutoModelForSeq2SeqLM, T5Tokenizer

# Load mT5-small for demonstration
model_name = "google/mt5-small"
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Move to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)
model_params = sum(p.numel() for p in model.parameters())

def demonstrate_span_corruption(text, language, tokenizer):
    """
    Simulate span corruption for demonstration.
    In practice, this happens during pre-training data preparation.
    """
    tokens = tokenizer.tokenize(text)

    # Select spans to corrupt (simplified: corrupt every 4th-7th token)
    corrupted = []
    targets = []
    sentinel_id = 0

    i = 0
    while i < len(tokens):
        if i > 0 and i % 5 == 0 and i + 2 < len(tokens):
            # Corrupt a span of 2-3 tokens
            span_len = min(3, len(tokens) - i)
            corrupted.append(f"<extra_id_{sentinel_id}>")
            targets.extend(
                [f"<extra_id_{sentinel_id}>"] + tokens[i : i + span_len]
            )
            sentinel_id += 1
            i += span_len
        else:
            corrupted.append(tokens[i])
            i += 1

    if sentinel_id > 0:
        targets.append(f"<extra_id_{sentinel_id}>")

    return " ".join(corrupted), " ".join(targets)


# Test on multiple languages
test_texts = {
    "English": "Natural language processing enables computers to understand human language.",
    "Spanish": "El procesamiento del lenguaje natural permite que las computadoras entiendan el idioma humano.",
    "German": "Die Verarbeitung natürlicher Sprache ermöglicht es Computern, menschliche Sprache zu verstehen.",
    "Japanese": "自然言語処理により、コンピュータは人間の言語を理解できるようになります。",
}

from torch.utils.data import Dataset


class MultilingualQADataset(Dataset):
    """
    Dataset for multilingual question answering.
    Formats data as text-to-text for mT5.
    """

    def __init__(
        self, examples, tokenizer, max_input_length=512, max_target_length=128
    ):
        self.examples = examples
        self.tokenizer = tokenizer
        self.max_input_length = max_input_length
        self.max_target_length = max_target_length

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

    def __getitem__(self, idx):
        example = self.examples[idx]

        # Format input as: "question: {question} context: {context}"
        input_text = (
            f"question: {example['question']} context: {example['context']}"
        )
        target_text = example["answer"]

        # Tokenize
        inputs = self.tokenizer(
            input_text,
            max_length=self.max_input_length,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        )

        targets = self.tokenizer(
            target_text,
            max_length=self.max_target_length,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        )

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


# Example multilingual QA data
multilingual_qa_examples = [
    {
        "question": "What is machine learning?",
        "context": "Machine learning is a subset of artificial intelligence that enables systems to learn from data.",
        "answer": "a subset of artificial intelligence",
    },
    {
        "question": "¿Qué es el aprendizaje automático?",
        "context": "El aprendizaje automático es un subconjunto de la inteligencia artificial que permite a los sistemas aprender de los datos.",
        "answer": "un subconjunto de la inteligencia artificial",
    },
    {
        "question": "Was ist maschinelles Lernen?",
        "context": "Maschinelles Lernen ist ein Teilgebiet der künstlichen Intelligenz, das Systemen ermöglicht, aus Daten zu lernen.",
        "answer": "ein Teilgebiet der künstlichen Intelligenz",
    },
]

def generate_text(model, tokenizer, prompt, max_length=50):
    """Generate text from a prompt using mT5."""
    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    outputs = model.generate(
        **inputs,
        max_length=max_length,
        num_beams=4,
        early_stopping=True,
        no_repeat_ngram_size=2,
    )

    return tokenizer.decode(outputs[0], skip_special_tokens=True)


# Test prompts in different languages
prompts = [
    "translate English to German: Hello, how are you?",
    "translate English to Spanish: The weather is nice today.",
    "summarize: Machine learning is a branch of artificial intelligence focused on building systems that learn from data.",
]