74 lines
1.9 KiB
Python
74 lines
1.9 KiB
Python
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)
|