import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np

class SurnameDataset(Dataset):
    def __init__(self, data):
        self.data = data
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]
        
# 假設(shè)我們的數(shù)據(jù)格式為 (surname, country)，例如 ('Smith', 'USA')
# 這里假設(shè)數(shù)據(jù)已經(jīng)預(yù)處理成對應(yīng)的數(shù)值表示
# 例如將字符映射為數(shù)字，國家名稱映射為數(shù)字等
 
# 數(shù)據(jù)預(yù)處理函數(shù)
def preprocess_data(data):
    processed_data = []
    for surname, country in data:
        # 將姓氏轉(zhuǎn)換為字符索引列表
        surname_indices = [char_to_index[char] for char in surname]
        # 將國家轉(zhuǎn)換為對應(yīng)的數(shù)字
        country_index = country_to_index[country]
        processed_data.append((surname_indices, country_index))
    return processed_data

class SurnameRNN(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SurnameRNN, self).__init__()
        self.hidden_size = hidden_size
        self.embedding = nn.Embedding(input_size, hidden_size)
        self.lstm = nn.LSTM(hidden_size, hidden_size)
        self.fc = nn.Linear(hidden_size, output_size)
 
    def forward(self, input, hidden):
        embedded = self.embedding(input).view(1, 1, -1)
        output, hidden = self.lstm(embedded, hidden)
        output = self.fc(output.view(1, -1))
        return output, hidden
 
    def init_hidden(self):
        return (torch.zeros(1, 1, self.hidden_size), torch.zeros(1, 1, self.hidden_size))

input_size = len(char_to_index)  # 姓氏中字符的數(shù)量
hidden_size = 128
output_size = len(country_to_index)  # 國家的數(shù)量
learning_rate = 0.001
num_epochs = 10
 
model = SurnameRNN(input_size, hidden_size, output_size)
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)
 
# 假設(shè)我們有一個經(jīng)過預(yù)處理的數(shù)據(jù)集 surname_data
# 數(shù)據(jù)格式為 (surname_indices, country_index)
 
# 將數(shù)據(jù)劃分為訓(xùn)練集和測試集
train_data = surname_data[:800]
test_data = surname_data[800:]
 
# 開始訓(xùn)練
for epoch in range(num_epochs):
    total_loss = 0
    for surname_indices, country_index in train_data:
        model.zero_grad()
        hidden = model.init_hidden()
        surname_tensor = torch.tensor(surname_indices, dtype=torch.long)
        country_tensor = torch.tensor([country_index], dtype=torch.long)
 
        for i in range(len(surname_indices)):
            output, hidden = model(surname_tensor[i], hidden)
        
        loss = criterion(output, country_tensor)
        total_loss += loss.item()
        loss.backward()
        optimizer.step()
    
    print('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, total_loss / len(train_data)))

# 導(dǎo)入必要的庫
import torch
from torch.utils.data import DataLoader
 
# 定義測試函數(shù)
def test_model(model, test_data):
    model.eval()  # 設(shè)置模型為評估模式
    correct = 0
    total = 0
    with torch.no_grad():
        for surname_indices, country_index in test_data:
            surname_tensor = torch.tensor(surname_indices, dtype=torch.long)
            country_tensor = torch.tensor([country_index], dtype=torch.long)
            
            hidden = model.init_hidden()
            
            for i in range(len(surname_indices)):
                output, hidden = model(surname_tensor[i], hidden)
            
            _, predicted = torch.max(output.data, 1)
            
            total += 1
            if predicted == country_tensor:
                correct += 1
    
    accuracy = correct / total
    print('Accuracy on test data: {:.2%}'.format(accuracy))
 
# 加載測試數(shù)據(jù)集
test_dataset = SurnameDataset(test_data)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=True)
 
# 加載已經(jīng)訓(xùn)練好的模型
model_path = "path_to_your_trained_model.pt"
model = torch.load(model_path)
 
# 測試模型
test_model(model, test_loader)