Upload modelling_RW.py

modelling_RW.py

@@ -276,9 +276,23 @@ class Attention(nn.Module):
             key_layer_ = key_layer.reshape(batch_size, self.num_kv, -1, self.head_dim)
             value_layer_ = value_layer.reshape(batch_size, self.num_kv, -1, self.head_dim)
 
-            attn_output = F.scaled_dot_product_attention(
-                query_layer_, key_layer_, value_layer_, None, 0.0, is_causal=True
-            )
+            if torch.__version__ < "2.0.0":
+                mask = torch.ones(q_length, q_length, device=query_layer_.device)
+                mask = torch.tril(mask)
+                mask = (1.0 - mask) * -10000
+                mask = mask.repeat(batch_size, 1, 1, 1)
+
+                scores = torch.matmul(query_layer_, key_layer_.transpose(-2, -1))
+                scores = scores / math.sqrt(float(self.head_dim))
+                scores = scores + mask.type_as(scores)
+
+                probs = nn.Softmax(dim=-1)(scores)
+
+                attn_output = probs @ value_layer_
+            else:
+                attn_output = F.scaled_dot_product_attention(
+                    query_layer_, key_layer_, value_layer_, None, 0.0, is_causal=True
+                )
 
             x = attn_output.view(batch_size, self.num_heads, q_length, self.head_dim)
             x = x.permute(0, 2, 1, 3)
@@ -945,7 +959,7 @@ class RWForTokenClassification(RWPreTrainedModel):
         else:
             classifier_dropout = 0.1
         self.dropout = nn.Dropout(classifier_dropout)
-        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        self.classifier = nn.Linear(config.hidden_size, config.num_lab els)
 
         # Initialize weights and apply final processing
         self.post_init()