исправь синтетические ошибки import os
import torch
import torch.nn as nn
import torch.optim as optim
from collections import Counter
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
from datasets import load_dataset # импортируем функцию для загрузки датасета

Задаем идентификатор датасета
dataset_id = "Jiayi-Pan/Countdown-Tasks-3to4"

Загружаем датасет, перемешиваем и выбираем первые 50000 примеров
dataset = load_dataset(dataset_id, split="train").shuffle(seed=42).select(range(50000))

Разбиваем датасет на обучающую и валидационную выборки (80/20)
split_dataset = dataset.train_test_split(test_size=0.2, seed=42)
train_data = split_dataset["train"]
val_data = split_dataset["test"]

1. Определение кастомного датасета для текстовых данных, работающего с данными в памяти
class CustomTextDataset(Dataset):
def init(self, data, vocab=None, build_vocab=False):
"""
data - объект датасета (например, список словарей) с полями "text" и "label".
Если build_vocab=True, то на основе данного датасета будет построен словарь.
Если vocab передаётся, он используется для преобразования токенов в индексы.
"""
self.samples = []
self.labels = []
self.max_length = 100 # максимальная длина последовательности (при необходимости можно изменить)

text
    for example in data:
        # Поля в датасете должны называться "text" и "label". При необходимости замените имена полей.
        text, label = example["text"], example["label"]
        # Простая токенизация по пробелам
        tokens = text.strip().split()
        self.samples.append(tokens)
        self.labels.append(int(label))
    
    if build_vocab:
        # Строим словарь: собираем все токены в датасете
        counter = Counter()
        for tokens in self.samples:
            counter.update(tokens)
        # Добавляем специальные токены
        self.vocab = {'<PAD>': 0, '<UNK>': 1}
        for token in counter:
            self.vocab[token] = len(self.vocab)
    else:
        # Если build_vocab == False, должен быть передан словарь
        if vocab is None:
            raise ValueError("Если build_vocab=False, необходимо передать словарь vocab")
        self.vocab = vocab

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

def __getitem__(self, idx):
    tokens = self.samples[idx]
    label = self.labels[idx]
    # Преобразуем токены в индексы, используя словарь. Если токена нет в словаре, используется <UNK>
    indices = [self.vocab.get(token, self.vocab['<UNK>']) for token in tokens]
    return torch.tensor(indices, dtype=torch.long), torch.tensor(label, dtype=torch.long)
Функция для формирования батчей с padding (выравнивание последовательностей до максимальной длины в батче)
def collate_fn(batch):
texts, labels = zip(*batch)
texts = pad_sequence(texts, batch_first=True, padding_value=0) # 0 соответствует <PAD>
labels = torch.stack(labels)
return texts, labels

2. Определение модели TextLSTM для классификации
class TextLSTM(nn.Module):
def init(self, vocab_size, embed_size, hidden_size, num_classes, num_layers=1, bidirectional=False):
super(TextLSTM, self).init()
self.embedding = nn.Embedding(vocab_size, embed_size, padding_idx=0)
self.lstm = nn.LSTM(embed_size, hidden_size, num_layers=num_layers,
batch_first=True, bidirectional=bidirectional)
lstm_output_size = hidden_size * 2 if bidirectional else hidden_size
self.fc = nn.Linear(lstm_output_size, num_classes)

text
def forward(self, x):
    # x: [batch_size, seq_len]
    embedded = self.embedding(x)  # [batch_size, seq_len, embed_size]
    lstm_out, (h_n, c_n) = self.lstm(embedded)
    if self.lstm.bidirectional:
        h_n = torch.cat((h_n[-2], h_n[-1]), dim=1)
    else:
        h_n = h_n[-1]
    out = self.fc(h_n)
    return out
3. Основной скрипт для обучения и валидации
def main():
# Создаем тренировочный датасет и строим словарь
train_dataset = CustomTextDataset(train_data, build_vocab=True)
vocab = train_dataset.vocab
vocab_size = len(vocab)

text
# Создаем валидационный датасет с использованием созданного словаря
val_dataset = CustomTextDataset(val_data, vocab=vocab, build_vocab=False)

batch_size = 32
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)

# Параметры модели
embed_size = 100
hidden_size = 128
num_classes = len(set(train_dataset.labels))
num_epochs = 10
learning_rate = 0.001

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = TextLSTM(vocab_size, embed_size, hidden_size, num_classes, bidirectional=True)
model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(num_epochs):
    model.train()
    total_loss = 0.0
    for texts, labels in train_loader:
        texts = texts.to(device)
        labels = labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(texts)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    
    avg_loss = total_loss / len(train_loader)
    print(f'Epoch {epoch+1}/{num_epochs}, Training Loss: {avg_loss:.4f}')
    
    # Валидация
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for texts, labels in val_loader:
            texts = texts.to(device)
            labels = labels.to(device)
            outputs = model(texts)
            _, predicted = torch.max(outputs, dim=1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    val_acc = correct / total if total > 0 else 0
    print(f'Validation Accuracy: {val_acc:.4f}')

# Сохранение обученной модели
torch.save(model.state_dict(), "text_lstm_model.pth")
print("Модель сохранена.")
if name == "main":
main()