Spaces:
Sleeping
Sleeping
File size: 4,114 Bytes
379ec43 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer, BertForSequenceClassification, AdamW, get_linear_schedule_with_warmup
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
import numpy as np
# Load the dataset
chat_transcripts = ["chat transcript 1", "chat transcript 2", ...]
survey_responses = [3.5, 4.2, ...] # Numerical survey responses
# Split the data into training, validation, and testing sets
train_texts, temp_texts, train_labels, temp_labels = train_test_split(chat_transcripts, survey_responses, test_size=0.3, random_state=42)
val_texts, test_texts, val_labels, test_labels = train_test_split(temp_texts, temp_labels, test_size=0.5, random_state=42)
# Pre-process the data using BERT tokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
train_encodings = tokenizer(train_texts, truncation=True, padding=True)
val_encodings = tokenizer(val_texts, truncation=True, padding=True)
test_encodings = tokenizer(test_texts, truncation=True, padding=True)
class SurveyDataset(Dataset):
def __init__(self, encodings, labels):
self.encodings = encodings
self.labels = labels
def __getitem__(self, idx):
item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
item['labels'] = torch.tensor(self.labels[idx], dtype=torch.float)
return item
def __len__(self):
return len(self.labels)
train_dataset = SurveyDataset(train_encodings, train_labels)
val_dataset = SurveyDataset(val_encodings, val_labels)
test_dataset = SurveyDataset(test_encodings, test_labels)
# Fine-tune the BERT model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=1).to(device)
train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=8, shuffle=False)
optim = AdamW(model.parameters(), lr=2e-5)
num_epochs = 3
num_training_steps = num_epochs * len(train_loader)
lr_scheduler = get_linear_schedule_with_warmup(optim, num_warmup_steps=0, num_training_steps=num_training_steps)
for epoch in range(num_epochs):
model.train()
for batch in train_loader:
optim.zero_grad()
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
labels = batch["labels"].unsqueeze(1).to(device)
outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
loss = outputs.loss
loss.backward()
optim.step()
lr_scheduler.step()
# Evaluate the model
model.eval()
preds = []
with torch.no_grad():
for batch in test_loader:
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
outputs = model(input_ids, attention_mask=attention_mask)
logits = outputs.logits
preds.extend(logits.squeeze().tolist())
mse = mean_squared_error(test_labels, preds)
r2 = r2_score(test_labels, preds)
print("Mean Squared Error:", mse)
print("R-squared Score:", r2)
def predict_survey_response(chat_transcript, model, tokenizer, device):
# Preprocess the chat transcript
encoding = tokenizer(chat_transcript, truncation=True, padding=True, return_tensors="pt")
# Move tensors to the device
input_ids = encoding["input_ids"].to(device)
attention_mask = encoding["attention_mask"].to(device)
# Predict the survey response using the fine-tuned model
model.eval()
with torch.no_grad():
outputs = model(input_ids, attention_mask=attention_mask)
logits = outputs.logits
predicted_response = logits.squeeze().item()
return predicted_response
# Example usage
new_chat_transcript = "A new chat transcript"
predicted_response = predict_survey_response(new_chat_transcript, model, tokenizer, device)
print("Predicted survey response:", predicted_response)
|