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(),
            "label": torch.tensor(label, dtype=torch.long)
        }

# 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_))
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),
            "label": batch["label"].to(device)
        }
        loss, _ = model(**inputs)
        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),
            "label": batch["label"].to(device)
            }
            _, outputs = model(**inputs)
            _, predicted = torch.max(outputs, 1)
            total_correct += (predicted == batch["label"].to(device)).sum().item()
            total_count += len(batch)
        
        accuracy = total_correct / total_count
        
        print(f"Epoch {epoch}, Accuracy: {accuracy}")   "label": batch["label"].to(device) TypeError: BertForSequenceClassification.forward() got an unexpected keyword argument 'label'