KameronB/SITCCSA

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))

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)