import os
import pandas as pd
from tqdm import tqdm
import torch
from torch.utils.data import DataLoader, TensorDataset
from transformers import (
    DistilBertTokenizerFast,
    DistilBertForSequenceClassification
)

# Importing the csv dataset.
case = '...'
df_data = pd.read_csv(f'{case}.csv')
print(len(df_data))

# Set device cuda
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using:', device)

# Define model
model_name = 'distilbert-base-uncased'
tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)

# Tokenize and encode the text data
encodings = tokenizer(
    df_data['text'].tolist(), 
    truncation=True, 
    padding=True, 
    return_tensors='pt'
)

# Create tensor dataset
dataset = TensorDataset(
    encodings['input_ids'], 
    encodings['attention_mask']
)

# Create data loader
if device == 'cuda':
    en_pin_memory = True
else:
    en_pin_memory = False
dataloader = DataLoader(
    dataset, 
    batch_size=53, 
    shuffle=False,
    pin_memory=en_pin_memory
)

# Loading the model
model = DistilBertForSequenceClassification.from_pretrained(model_name, num_labels=1)
model.load_state_dict(torch.load('...', map_location=device))
model.to(device)

# Using the model to perform sentiment analysis on the dataset.
model.eval()
list_predicted_scores = []

for batch in tqdm(dataloader):
    with torch.no_grad():
        input_ids, attention_mask = [x.to(device) for x in batch]

        # Obtaining the sentiment score.
        output = model(input_ids=input_ids, attention_mask=attention_mask)
        predicted_scores = output.logits.view(-1)

        # Writing the sentiment score to the list.
        list_predicted_scores.extend(predicted_scores.tolist())

# Inserting the sentiment score in the dataset.
df_data['s'] = list_predicted_scores

# Exporting the dataframe to a csv dataset.
df_data.to_csv(f'{case}-s.csv', index=False)