import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder
import torch
from torch.utils.data import DataLoader, Dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AdamW

# define the model repo
model_name = "cointegrated/rubert-tiny"

# load the dataset
#df = pd.read_csv(‘dataset.csv’)

# preprocess the text data
tokenizer = AutoTokenizer.from_pretrained(model_name)

class CustomDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length
    
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.labels[idx]
        
        # tokenize the text using the tokenizer
        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            pad_to_max_length=True,
            return_tensors="pt"
        )
        
        return {
            "input_ids": encoding["input_ids"].squeeze(),
            "attention_mask": encoding["attention_mask"].squeeze(),
            "labels": torch.tensor(label, dtype=torch.long).to(device)
        }


device = "cuda" if torch.cuda.is_available() else "cpu"
# encode the labels
le = LabelEncoder()
train_data["label"] = le.fit_transform(train_data["category"])

# split the dataset into train and test set
X_train, X_test, y_train, y_test = train_test_split(train_data["text"], train_data["label"], test_size=0.2, random_state=42)

# create the dataloaders
train_dataset = CustomDataset(X_train.values, y_train.values, tokenizer, max_length=512)
test_dataset = CustomDataset(X_test.values, y_test.values, tokenizer, max_length=512)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)


# define the model and its optimizer
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=len(le.classes_)).to(device)
optimizer = AdamW(model.parameters(), lr=2e-5)

# training loop
num_epochs = 5
for epoch in range(num_epochs):
    model.train()
    for batch in train_loader:
        inputs = {
            "input_ids": batch["input_ids"].to(device),
            "attention_mask": batch["attention_mask"].to(device),
            "labels": batch["labels"].to(device)
        }
        #loss, _ = model(**inputs)
        outputs = model(**inputs)
        loss = outputs.loss
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()
    
    model.eval()
    with torch.no_grad():
        total_correct = 0
        total_count = 0
        for batch in test_loader:

            inputs = {
            "input_ids": batch["input_ids"].to(device),
            "attention_mask": batch["attention_mask"].to(device),
            "labels": batch["labels"].to(device)
            }

            outputs = model(**inputs)
            _, predicted = torch.max(outputs.logits, 1)
            total_correct += (predicted == batch["labels"].to(device)).sum().item()
            total_count += len(batch)
        
        accuracy = total_correct / total_count
        
        print(f"Epoch {epoch}, Accuracy: {accuracy}") Я обучил таким образом модель. Затем я пытаюсь ее запустить max_length = 512

# получить список токенизированных последовательностей списков
tokenized = [tokenizer.encode(text, add_special_tokens=True, max_length=max_length) for text in text_list]

# определить максимальную длину последовательности в списке токенов
max_len = max([len(seq) for seq in tokenized])

# заполнить последовательности до максимальной длины и преобразовать в формат, который можно передать в модель
input_ids = torch.tensor([t + [0]*(max_len-len(t)) for t in tokenized], dtype=torch.long)
attention_mask = torch.tensor([[float(i > 0) for i in ii] for ii in input_ids], dtype=torch.float)

# передать токенизированный вектор input_ids и маску внимания attention_mask для получения
# векторных представлений для каждого токена
with torch.no_grad():
    outputs = model(input_ids=input_ids, attention_mask=attention_mask)
    embeddings = outputs.last_hidden_state

# среднее каждого векторного представления по всем токенам в слове
# для получения векторного представления всей последовательности
mean_embeddings = torch.mean(embeddings, dim=1)

# преобразовать векторные представления в массив numpy и проверить размерность
mean_embeddings = mean_embeddings.numpy()
print(mean_embeddings.shape)  RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper_CUDA__index_select)