Create daedalus_mobile.py

daedalus_mobile.py

@@ -0,0 +1,34 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+class DaedalusMobile(nn.Module):
+    def __init__(self, config):
+        super(DaedalusMobile, self).__init__()
+        self.config = config
+        self.encoder = AutoModelForSeq2SeqLM.from_pretrained('t5-small')
+        self.decoder = AutoModelForSeq2SeqLM.from_pretrained('t5-small')
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, input_ids, attention_mask):
+        encoder_output = self.encoder(input_ids, attention_mask)
+        decoder_output = self.decoder(encoder_output.last_hidden_state, attention_mask)
+        output = self.dropout(decoder_output.last_hidden_state)
+        return output
+
+    def configure_optimizers(self):
+        optimizer = optim.Adam(self.parameters(), lr=self.config.lr)
+        return optimizer
+
+    def train_step(self, batch):
+        input_ids, attention_mask, labels = batch
+        output = self(input_ids, attention_mask)
+        loss = nn.CrossEntropyLoss()(output, labels)
+        return loss
+
+    def eval_step(self, batch):
+        input_ids, attention_mask, labels = batch
+        output = self(input_ids, attention_mask)
+        loss = nn.CrossEntropyLoss()(output, labels)
+        return loss