Spaces:
Build error
Build error
File size: 5,484 Bytes
a446b0b 0935f1b a446b0b 0935f1b a446b0b |
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 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 |
import torch, argparse
from commonsense_model import CommonsenseGRUModel
from dataloader import RobertaCometDataset
from torch.utils.data import DataLoader
def load_model(model_path, args):
emo_gru = True
n_classes = 15
cuda = args.cuda
D_m = 1024
D_s = 768
D_g = 150
D_p = 150
D_r = 150
D_i = 150
D_h = 100
D_a = 100
D_e = D_p + D_r + D_i
model = CommonsenseGRUModel(
D_m,
D_s,
D_g,
D_p,
D_r,
D_i,
D_e,
D_h,
D_a,
n_classes=n_classes,
listener_state=args.active_listener,
context_attention=args.attention,
dropout_rec=args.rec_dropout,
dropout=args.dropout,
emo_gru=emo_gru,
mode1=args.mode1,
norm=args.norm,
residual=args.residual,
)
if cuda:
model.cuda()
model.load_state_dict(torch.load(model_path))
else:
model.load_state_dict(torch.load(model_path, map_location=torch.device("cpu")))
model.eval()
return model
def get_valid_dataloader(
roberta_features_path: str,
comet_features_path: str,
batch_size=1,
num_workers=0,
pin_memory=False,
):
valid_set = RobertaCometDataset("valid", roberta_features_path, comet_features_path)
test_loader = DataLoader(
valid_set,
batch_size=batch_size,
collate_fn=valid_set.collate_fn,
num_workers=num_workers,
pin_memory=pin_memory,
)
return test_loader, valid_set.keys
def predict(model, data_loader, args):
predictions = []
for data in data_loader:
r1, r2, r3, r4, x1, x2, x3, x4, x5, x6, o1, o2, o3, qmask, umask, label = (
[d.cuda() for d in data[:-1]] if args.cuda else data[:-1]
)
log_prob, _, alpha, alpha_f, alpha_b, _ = model(
r1, r2, r3, r4, x5, x6, x1, o2, o3, qmask, umask
)
lp_ = log_prob.transpose(0, 1).contiguous().view(-1, log_prob.size()[2])
preds = torch.argmax(lp_, dim=-1)
predictions.append(preds.data.cpu().numpy())
return predictions
def parse_cosmic_args():
parser = argparse.ArgumentParser()
# Parse arguments input into the cosmic model
parser.add_argument(
"--no-cuda", action="store_true", default=True, help="does not use GPU"
)
parser.add_argument(
"--lr", type=float, default=0.0001, metavar="LR", help="learning rate"
)
parser.add_argument(
"--l2",
type=float,
default=0.00003,
metavar="L2",
help="L2 regularization weight",
)
parser.add_argument(
"--rec-dropout",
type=float,
default=0.3,
metavar="rec_dropout",
help="rec_dropout rate",
)
parser.add_argument(
"--dropout", type=float, default=0.5, metavar="dropout", help="dropout rate"
)
parser.add_argument(
"--batch-size", type=int, default=1, metavar="BS", help="batch size"
)
parser.add_argument(
"--epochs", type=int, default=10, metavar="E", help="number of epochs"
)
parser.add_argument(
"--class-weight", action="store_true", default=True, help="use class weights"
)
parser.add_argument(
"--active-listener", action="store_true", default=True, help="active listener"
)
parser.add_argument(
"--attention", default="simple", help="Attention type in context GRU"
)
parser.add_argument(
"--tensorboard",
action="store_true",
default=False,
help="Enables tensorboard log",
)
parser.add_argument("--mode1", type=int, default=2, help="Roberta features to use")
parser.add_argument("--seed", type=int, default=500, metavar="seed", help="seed")
parser.add_argument("--norm", type=int, default=0, help="normalization strategy")
parser.add_argument("--mu", type=float, default=0, help="class_weight_mu")
parser.add_argument(
"--residual", action="store_true", default=True, help="use residual connection"
)
args = parser.parse_args()
args.cuda = torch.cuda.is_available() and not args.no_cuda
if args.cuda:
print("Running on GPU")
else:
print("Running on CPU")
return args
if __name__ == "__main__":
def pred_to_labels(preds):
mapped_predictions = []
for pred in preds:
# map the prediction for each conversation
mapped_labels = []
for label in pred:
mapped_labels.append(label_mapping[label])
mapped_predictions.append(mapped_labels)
# return the mapped labels for each conversation
return mapped_predictions
label_mapping = {
0: "Curiosity",
1: "Obscene",
2: "Informative",
3: "Openness",
4: "Acceptance",
5: "Interest",
6: "Greeting",
7: "Disapproval",
8: "Denial",
9: "Anxious",
10: "Uninterested",
11: "Remorse",
12: "Confused",
13: "Accusatory",
14: "Annoyed",
}
args = parse_cosmic_args()
model = load_model("epik/best_model.pt", args)
test_dataloader, ids = get_valid_dataloader()
predicted_labels = pred_to_labels(predict(model, test_dataloader, args))
for id, labels in zip(ids, predicted_labels):
print(f"Conversation ID: {id}")
print(f"Predicted Sentiment Labels: {labels}")
print(len(labels))
|