|
|
--- |
|
|
license: mit |
|
|
--- |
|
|
|
|
|
|
|
|
```python |
|
|
import torch |
|
|
from torch import nn |
|
|
from transformers import RobertaConfig, RobertaModel, RobertaTokenizer |
|
|
import torch.optim as optim |
|
|
|
|
|
|
|
|
class HelpdeskReviewModel(nn.Module): |
|
|
def __init__(self): |
|
|
super(HelpdeskReviewModel, self).__init__() |
|
|
self.bert = RobertaModel.from_pretrained('roberta-base') |
|
|
self.drop = nn.Dropout(p=0.2) |
|
|
|
|
|
# Fully connected layers |
|
|
self.fc1 = nn.Linear(self.bert.config.hidden_size, 512) # First fully connected layer |
|
|
self.fc2 = nn.Linear(512, 256) # Second fully connected layer |
|
|
self.fc3 = nn.Linear(256, 128) # Third fully connected layer |
|
|
|
|
|
# Activation function |
|
|
self.relu = nn.ReLU() |
|
|
|
|
|
self.output = nn.Linear(128, 4) # 4 outputs |
|
|
self.sigmoid = nn.Sigmoid() |
|
|
|
|
|
state_dict = torch.hub.load_state_dict_from_url(f"https://huggingface.co/KameronB/SITCCSA/resolve/main/pytorch_model.bin", model_dir="./model", file_name="sitccsa_pytorch_model.bin") |
|
|
# if running on cpu |
|
|
# state_dict = torch.hub.load_state_dict_from_url(f"https://huggingface.co/KameronB/SITCCSA/resolve/main/pytorch_model.bin", model_dir="./model", file_name="sitccsa_pytorch_model.bin", map_location=torch.device('cpu')) |
|
|
|
|
|
|
|
|
self.load_state_dict(state_dict) |
|
|
|
|
|
def forward(self, input_ids, attention_mask): |
|
|
_, pooled_output = self.bert( |
|
|
input_ids=input_ids, |
|
|
attention_mask=attention_mask, |
|
|
return_dict=False |
|
|
) |
|
|
output = self.drop(pooled_output) |
|
|
|
|
|
# apply new hidden layers |
|
|
output = self.relu(self.fc1(output)) |
|
|
output = self.relu(self.fc2(output)) |
|
|
output = self.relu(self.fc3(output)) |
|
|
|
|
|
return self.sigmoid(self.output(output)) |
|
|
``` |
|
|
|
|
|
```python |
|
|
|
|
|
tokenizer = RobertaTokenizer.from_pretrained('roberta-base') |
|
|
model = HelpdeskReviewModel() |
|
|
|
|
|
|
|
|
def make_prediction(model, tokenizer, text, max_length=512): |
|
|
# Tokenize the input text and convert to tensor |
|
|
inputs = tokenizer.encode_plus( |
|
|
text, |
|
|
add_special_tokens=True, # Add '[CLS]' and '[SEP]' |
|
|
max_length=max_length, # Pad & truncate all sentences. |
|
|
padding='max_length', |
|
|
truncation=True, |
|
|
return_tensors="pt" # Return PyTorch tensors. |
|
|
) |
|
|
|
|
|
# Move tensors to the same device as model |
|
|
input_ids = inputs['input_ids'] |
|
|
attention_mask = inputs['attention_mask'] |
|
|
if torch.cuda.is_available(): |
|
|
input_ids = input_ids.cuda() |
|
|
attention_mask = attention_mask.cuda() |
|
|
model = model.cuda() |
|
|
|
|
|
# Make prediction |
|
|
with torch.no_grad(): |
|
|
outputs = model(input_ids, attention_mask=attention_mask) |
|
|
|
|
|
# Return probabilities |
|
|
return outputs.cpu().numpy()[0] # Return to CPU and convert to numpy array if not running on CPU |
|
|
|
|
|
# Example usage of the function |
|
|
texts = [ |
|
|
"This Agent is TERRIBLE!: The agent I spoke to on the phone did not seem to have any idea of what he was doing.", |
|
|
"Excellent work!: The tech that installed my software was amazing! Thank you!", |
|
|
"Good Work: The person who anwsered the phone did a pretty good job. It took a bit longer than I would have liked, but they got the job done.", |
|
|
"Bad Computer: My Computer is a piece of junk!!!", |
|
|
"Poor Service: I sent David and email and it took him over 30 seconds to respond. The service is so slow that I missed the solar eclipse.", |
|
|
"Very Slow: The technician was very slow.", |
|
|
"Thank you!: Stanley did a great job installing my software!", |
|
|
"You need better training: These agents need better training, they cant even seem to do simple troubleshooting.", |
|
|
"The technician threatened my life: The technician threatened my life!" |
|
|
] |
|
|
for text in texts: |
|
|
probabilities = make_prediction(model, tokenizer, text) |
|
|
print(probabilities) |
|
|
``` |
