У меня есть код для получения эмбеддингов текста и решения задачи многоклассовой классификации на 50 классов, однако качество меня не устраивает. Я хочу использовать архитектуру cross-encoder с помощью модели cross-encoder/ms-marco-MiniLM-L-12-v2 которая будет брать текст и сравнивать его с 50 описаниями из переменной categories. Вот текущий код обработки:

model_name = 'intfloat/multilingual-e5-large'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
model = torch.nn.DataParallel(model, device_ids = [0,1]).to(device)
model.eval()

text_embeddings_train_np_2 = embed_text(train['text'].tolist())

print("Generating text embeddings (test data)...")
text_embeddings_test_np_2 = embed_text(test['text'].tolist())

target_cols = [f'trend_id_res{i}' for i in range(50)]
y_train = train[target_cols].values

X_train = np.hstack((tags_encoded, text_embeddings_train_np_2))
X_test = np.hstack((tags_encoded_test, text_embeddings_test_np_2,))

X_tr, X_val, y_tr, y_val = train_test_split(X_train, y_train, test_size=0.2, random_state=42)

class CustomDataset(Dataset):
    def __init__(self, inputs, targets):
        self.inputs = torch.tensor(inputs, dtype=torch.float32)
        self.targets = torch.tensor(targets, dtype=torch.float32)
        
    def __len__(self):
        return len(self.inputs)
    
    def __getitem__(self, index):
        x = self.inputs[index]
        y = self.targets[index]
        return x, y
    
batch_size = 64

train_dataset = CustomDataset(X_tr, y_tr)
val_dataset = CustomDataset(X_val, y_val)

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

class MultiLabelNN(nn.Module):
    def __init__(self, input_size, hidden_sizes, output_size):
        super(MultiLabelNN, self).__init__()
        layers = []
        last_size = input_size
        for hidden_size in hidden_sizes:
            layers.append(nn.Linear(last_size, hidden_size))
            layers.append(nn.BatchNorm1d(hidden_size))
            layers.append(nn.LeakyReLU(0.1))
            layers.append(nn.Dropout(0.2))
            last_size = hidden_size
        layers.append(nn.Linear(last_size, output_size))
        layers.append(nn.Sigmoid())
        self.model = nn.Sequential(*layers)
        
    def forward(self, x):
        return self.model(x)
    
input_size = X_tr.shape[1]
hidden_sizes = [1024, 512]
output_size = y_tr.shape[1]

model_nn = MultiLabelNN(input_size, hidden_sizes, output_size)
model_nn.to(device)

criterion = nn.BCELoss()
optimizer = optim.AdamW(model_nn.parameters(), lr=3e-4)

def exact_match(y_true, y_pred):
    y_pred_labels = (y_pred > 0.5).float()
    matches = (y_pred_labels == y_true).all(dim=1)
    return matches.float().mean().item()

num_epochs = 20

for epoch in range(num_epochs):
    model_nn.train()
    total_loss = 0
    for inputs, targets in tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}"):
        inputs = inputs.to(device)
        targets = targets.to(device)
        
        optimizer.zero_grad()
        outputs = model_nn(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        
        total_loss += loss.item()
    
    avg_loss = total_loss / len(train_loader)
    
    # Validation
    model_nn.eval()
    val_outputs = []
    val_targets = []
    with torch.no_grad():
        for inputs, targets in val_loader:
            inputs = inputs.to(device)
            targets = targets.to(device)
            outputs = model_nn(inputs)
            val_outputs.append(outputs.cpu())
            val_targets.append(targets.cpu())
    val_outputs = torch.cat(val_outputs, dim=0)
    val_targets = torch.cat(val_targets, dim=0)
    val_accuracy = exact_match(val_targets, val_outputs)
    
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.4f}, Validation Accuracy (Exact Match): {val_accuracy:.4f}")

from torch.utils.data import TensorDataset
X_test_tensor = torch.tensor(X_test, dtype=torch.float32)

# Создаем DataLoader для тестовых данных
test_dataset = TensorDataset(X_test_tensor)
test_loader = DataLoader(test_dataset, batch_size=batch_size)

# Генерация предсказаний на тестовых данных
model_nn.eval()
test_outputs = []
with torch.no_grad():
    for inputs in tqdm(test_loader, desc="Предсказание на тестовых данных"):
        inputs = inputs[0].to(device)
        outputs = model_nn(inputs)
        test_outputs.append(outputs.cpu())
test_outputs = torch.cat(test_outputs, dim=0)

test_pred_labels = (test_outputs > 0.5).int().numpy()

# Создаем датафрейм с результатами
pred_df = pd.DataFrame(test_pred_labels, columns=target_cols)
pred_df.insert(0, 'index', test['index'].values)

# Сохраняем в CSV
pred_df