import torch
import torch.nn as nn

class Encoder(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)

    def forward(self, x):
        embedded = self.embedding(x)
        outputs, (hidden, cell) = self.lstm(embedded)
        return hidden, cell

class Decoder(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)

    def forward(self, input, hidden, cell):
        embedded = self.embedding(input).unsqueeze(1)
        output, (hidden, cell) = self.lstm(embedded, (hidden, cell))
        prediction = self.fc(output.squeeze(1))
        return prediction, hidden, cell

class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, src, trg, teacher_forcing_ratio=0.5):
        batch_size = src.shape[0]
        trg_len = trg.shape[1]
        trg_vocab_size = self.decoder.fc.out_features

        outputs = torch.zeros(batch_size, trg_len, trg_vocab_size)
        hidden, cell = self.encoder(src)

        input = trg[:, 0]  # <sos> token

        for t in range(1, trg_len):
            output, hidden, cell = self.decoder(input, hidden, cell)
            outputs[:, t] = output

            teacher_force = torch.rand(1).item() < teacher_forcing_ratio
            top1 = output.argmax(1)
            input = trg[:, t] if teacher_force else top1

        return outputs

class Attention(nn.Module):
    def __init__(self, hidden_dim):
        super().__init__()
        self.attn = nn.Linear(hidden_dim * 2, hidden_dim)
        self.v = nn.Linear(hidden_dim, 1, bias=False)

    def forward(self, hidden, encoder_outputs):
        # hidden: (batch, hidden_dim)
        # encoder_outputs: (batch, src_len, hidden_dim)
        src_len = encoder_outputs.shape[1]
        hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)
        energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim=2)))
        attention = self.v(energy).squeeze(2)
        return torch.softmax(attention, dim=1)

from nltk.translate.bleu_score import corpus_bleu, sentence_bleu

# 句子级 BLEU
reference = [["the", "cat", "is", "on", "the", "mat"]]
hypothesis = ["the", "cat", "sat", "on", "the", "mat"]
score = sentence_bleu(reference, hypothesis, weights=(0.25, 0.25, 0.25, 0.25))
print(f"BLEU: {score:.4f}")

# 语料级 BLEU
references = [
    [["the", "cat", "is", "on", "the", "mat"], ["there", "is", "a", "cat", "on", "the", "mat"]],
    [["the", "dog", "is", "in", "the", "garden"]],
]
hypotheses = [
    ["the", "cat", "sat", "on", "the", "mat"],
    ["the", "dog", "is", "in", "the", "garden"],
]
score = corpus_bleu(references, hypotheses)
print(f"Corpus BLEU: {score:.4f}")

# 使用 GPT 进行翻译
prompt = """Translate the following English text to Chinese:

English: Natural language processing is a fascinating field.
Chinese:"""

# GPT-4 会生成："自然语言处理是一个引人入胜的领域。"

from transformers import MarianMTModel, MarianTokenizer

# 加载英中翻译模型
model_name = "Helsinki-NLP/opus-mt-en-zh"
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)

def translate(text, source_lang="en", target_lang="zh"):
    inputs = tokenizer(text, return_tensors="pt", padding=True)
    outputs = model.generate(**inputs, max_length=128)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# 测试
text = "Natural language processing is a fascinating field of computer science."
translation = translate(text)
print(f"English: {text}")
print(f"Chinese: {translation}")

import torch
import torch.nn as nn
import torch.optim as optim

# 超参数
SRC_VOCAB_SIZE = 5000
TRG_VOCAB_SIZE = 5000
EMBED_DIM = 256
HIDDEN_DIM = 512
N_LAYERS = 2

# 创建模型
encoder = Encoder(SRC_VOCAB_SIZE, EMBED_DIM, HIDDEN_DIM)
decoder = Decoder(TRG_VOCAB_SIZE, EMBED_DIM, HIDDEN_DIM)
model = Seq2Seq(encoder, decoder).to(device)

optimizer = optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)

# 训练循环
for epoch in range(N_EPOCHS):
    model.train()
    total_loss = 0

    for batch in train_loader:
        src = batch.src.to(device)
        trg = batch.trg.to(device)

        optimizer.zero_grad()
        output = model(src, trg)
        output = output[:, 1:].reshape(-1, TRG_VOCAB_SIZE)
        trg = trg[:, 1:].reshape(-1)

        loss = criterion(output, trg)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()

    print(f"Epoch {epoch}, Loss: {total_loss / len(train_loader):.4f}")

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

# 使用预训练 Transformer 翻译模型
model_name = "t5-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def translate_with_t5(text, target_lang="Chinese"):
    input_text = f"translate English to {target_lang}: {text}"
    inputs = tokenizer(input_text, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=128)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

text = "The weather is beautiful today."
print(translate_with_t5(text))