fine_tuned_model / bert.py
AnnaLissa's picture
Upload bert.py
ed98c3f verified
import pandas as pd
from sklearn.model_selection import train_test_split
from transformers import BertTokenizer, TFBertForSequenceClassification
import tensorflow as tf
# Load preprocessed data
def load_data(file_path="preprocessed_reviews.csv"):
return pd.read_csv("preprocessed_reviews.csv")
# Tokenize text using BERT tokenizer
def tokenize_text(tokenizer, texts, max_length):
encodings = tokenizer(texts.tolist(), padding=True, truncation=True, max_length=max_length, return_tensors="tf")
# Convert BatchEncoding to dictionary
encodings_dict = {key: value.numpy() for key, value in encodings.items()}
return encodings_dict
if __name__ == "__main__":
# Load preprocessed data
data = load_data("preprocessed_reviews.csv")
# Check if 'sentiment' column exists
if 'sentiment' in data.columns:
# Split data into train and validation sets
train_data, val_data = train_test_split(data, test_size=0.2, random_state=42)
# Tokenize text using BERT tokenizer
max_length = 128
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
train_inputs = tokenize_text(tokenizer, train_data['clean_text'], max_length)
val_inputs = tokenize_text(tokenizer, val_data['clean_text'], max_length)
# Convert 'sentiment' column to numerical format
num_labels = len(data['sentiment'].unique())
train_labels = train_data['sentiment'].astype('category').cat.codes.values
val_labels = val_data['sentiment'].astype('category').cat.codes.values
# Fine-tuning BERT model for sequence classification
model = TFBertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=num_labels)
optimizer = tf.keras.optimizers.Adam(learning_rate=2e-5)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
metrics = ['accuracy']
model.compile(optimizer=optimizer, loss=loss, metrics=metrics)
# Train the model
history = model.fit(
train_inputs,
train_labels,
validation_data=(val_inputs, val_labels),
epochs=3,
batch_size=32,
verbose=1
)
# Evaluate the model
loss, accuracy = model.evaluate(val_inputs, val_labels)
print(f'Validation loss: {loss}, Validation accuracy: {accuracy}')
# Save the trained model
model.save_pretrained('fine_tuned_bert_model')
else:
raise ValueError("The 'sentiment' column is not found in the DataFrame.")