The model predicts one of 3 classes. The mapping from ID to label is as follows:
1{
2 "0": "HUMAN_GENERATED",
3 "1": "MACHINE_GENERATED",
4 "2": "MACHINE_REFINED",
5}
1TEXT_EMBEDDING_DIM = 768
2
3
4class TLModel(nn.Module):
5 def __init__(self, text_encoder, projection_dim=128, num_classes=NUM_CLASSES, class_weights=None):
6 super().__init__()
7 self.text_encoder = text_encoder
8 self.num_classes = num_classes
9 text_output_dim = TEXT_EMBEDDING_DIM
10 self.additional_loss = losses.BatchHardSoftMarginTripletLoss(self.text_encoder)
11
12 self.text_projection = nn.Linear(text_output_dim, projection_dim)
13 self.classifier = nn.Linear(projection_dim, num_classes)
14 self.class_weights = class_weights
15
16 def forward(self, labels=None, input_ids=None, attention_mask=None):
17 actual_labels = labels
18 sentence_embeddings = self.text_encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
19 sentence_embeddings = sentence_embeddings.mean(dim=1)
20 projected_text = F.relu(self.text_projection(sentence_embeddings))
21 logits = self.classifier(projected_text)
22 loss = None
23 cross_entropy_loss = None
24 contrastive_loss = None
25
26 if actual_labels is not None:
27 loss_fct_ce = nn.CrossEntropyLoss(weight=self.class_weights.to(logits.device) if self.class_weights is not None else None)
28 cross_entropy_loss = loss_fct_ce(logits.view(-1, self.num_classes), actual_labels.view(-1))
29 contrastive_loss = self.additional_loss.batch_hard_triplet_loss(embeddings=projected_text, labels=actual_labels)
30 lambda_contrast = 0.1
31 loss = cross_entropy_loss + lambda_contrast * contrastive_loss
32
33
34 output = {"logits": logits, "fused_embedding": projected_text}
35 if loss is not None:
36 output["loss"] = loss
37 if cross_entropy_loss is not None:
38 output["cross_entropy_loss"] = cross_entropy_loss
39 if contrastive_loss is not None:
40 output["contrastive_loss"] = contrastive_loss
41
42 return output
